using resnet instead of efficientnet

train.py

@@ -3,6 +3,7 @@ from datasets import load_dataset
 import evaluate
 from transformers import EfficientNetImageProcessor, EfficientNetForImageClassification, TrainingArguments, Trainer
 import numpy as np
+from torchvision import models, transforms
 
 print("Cuda availability:", torch.cuda.is_available())
 cuda = torch.device('cuda')
@@ -20,10 +21,23 @@ for i, label in enumerate(labels):
     label2id[label] = str(i)
     id2label[str(i)] = label
 
-preprocessor = EfficientNetImageProcessor.from_pretrained(model_name)
-model = EfficientNetForImageClassification.from_pretrained(model_name, num_labels=len(
-    labels), id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True)
-
+# preprocessor = EfficientNetImageProcessor.from_pretrained(model_name)
+# model = EfficientNetForImageClassification.from_pretrained(model_name, num_labels=len(
+#     labels), id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True)
+    
+
+# Replace the EfficientNetImageProcessor with torchvision transforms
+preprocessor = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Replace the EfficientNetForImageClassification with torchvision ResNet-50
+model = models.resnet50(pretrained=True)
+num_ftrs = model.fc.in_features
+model.fc = torch.nn.Linear(num_ftrs, len(labels))
 
 training_args = TrainingArguments(
     finetuned_model_name, remove_unused_columns=False,