Update vit_model_training.py

vit_model_training.py

@@ -8,7 +8,7 @@ import torch.optim as optim
 import os
 import pandas as pd
 from sklearn.model_selection import train_test_split
-## working 18.5.24
+## working 18.9.24
 
 
 def labeling(path_real, path_fake):
@@ -36,13 +36,13 @@ class CustomDataset(Dataset):
         return len(self.dataframe)
 
     def __getitem__(self, idx):
-        image_path = self.dataframe.iloc[idx, 0]  # Image path is in the first column
+        image_path = self.dataframe.iloc[idx, 0]  
         image = Image.open(image_path).convert('RGB')  # Convert to RGB format
 
         if self.transform:
             image = self.transform(image)
 
-        label = self.dataframe.iloc[idx, 1]  # Label is in the second column
+        label = self.dataframe.iloc[idx, 1]  
         return image, label
     
     
@@ -62,21 +62,21 @@ if __name__ == "__main__":
     # Check for GPU availability
     device = torch.device('cuda')
 
-    # Load the pre-trained ViT model and move it to GPU
+    # Load the pre-trained ViT model
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(device)
 
     # Freeze pre-trained layers
     for param in model.parameters():
         param.requires_grad = False
 
-    # Define a new classifier and move it to GPU
+    # Define a new classifier
     model.classifier = nn.Linear(model.config.hidden_size, 2).to(device)  # Two output classes: 'REAL' and 'FAKE'
     
     print(model)
     # Define the optimizer
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-    # Define the image preprocessing pipeline
+    # Resize image and make it a tensor (add dimension)
     preprocess = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor()