Update train_mlp.py

train_mlp.py

@@ -1,15 +1,13 @@
-#imports
-
 import os
 import csv
 import torch
 from torch import nn 
 from torch.utils.data import DataLoader 
 from torchvision import datasets 
-from torchvision.transforms import ToTensor, Normalize, RandomCrop, RandomHorizontalFlip, Compose 
+from torchvision.transforms import ToTensor, Normalize, Compose 
 from contextualizer_mlp import Contextualizer
 
-# data transforms
+
 
 transform = Compose([
  
@@ -31,7 +29,7 @@ test_data = datasets.CIFAR10(
                                        download=True,
                                        transform=transform 
                                        )                                       
-# create dataloaders  
+
                                      
 batch_size = 128
 
@@ -44,22 +42,12 @@ for X, y in test_dataloader:
     print(f"Shape of y:{y.shape}{y.dtype}")
     break
 
-# size checking for loading images
-def check_sizes(image_size, patch_size):
-    sqrt_num_patches, remainder = divmod(image_size, patch_size)
-    assert remainder == 0, "`image_size` must be divisibe by `patch_size`"
-    num_patches = sqrt_num_patches ** 2
-    return num_patches
-
-
 
-# create model
-# Get cpu or gpu device for training.
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print(f"using {device} device") 
 
-# model definition
+
 
 class ContextualizerImageClassification(Contextualizer):
     def __init__(
@@ -74,7 +62,6 @@ class ContextualizerImageClassification(Contextualizer):
         num_layers=4,
         dropout=0.5
     ):
-        num_patches = check_sizes(image_size, patch_size)
         super().__init__(d_model,d_ffn,num_layers,dropout, num_tokens)
         self.patcher = nn.Conv2d(
             in_channels, d_model, kernel_size=patch_size, stride=patch_size
@@ -88,20 +75,20 @@ class ContextualizerImageClassification(Contextualizer):
         patches = patches.permute(0, 2, 3, 1)
         patches = patches.view(batch_size, -1, num_channels)
         embedding = self.model(patches)
-        embedding = embedding.mean(dim=1) # global average pooling
+        embedding = embedding.mean(dim=1) 
         out = self.classifier(embedding)
         return out
 
 model = ContextualizerImageClassification().to(device)
 print(model)
 
-# Optimizer
+
 
 loss_fn = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters(),lr=1e-3)
 
 
-# Training Loop
+
 
 def train(dataloader, model, loss_fn, optimizer):
     size = len(dataloader.dataset)
@@ -112,11 +99,11 @@ def train(dataloader, model, loss_fn, optimizer):
     for batch, (X,y) in enumerate(dataloader):
         X, y = X.to(device), y.to(device)
        
-        #compute prediction error
+      
         pred = model(X)
         loss = loss_fn(pred,y)
         
-        # backpropagation
+      
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
@@ -138,7 +125,6 @@ def train(dataloader, model, loss_fn, optimizer):
 
 
 
-# Test loop
 
 def test(dataloader, model, loss_fn):
     size = len(dataloader.dataset)            
@@ -160,7 +146,7 @@ def test(dataloader, model, loss_fn):
 
 
 
-# apply train and test
+
 
 logname = "/PATH/Contextualizer_mlp/Experiments_cifar10/logs_contextualizer/logs_cifar10.csv"
 if not os.path.exists(logname):
@@ -184,7 +170,7 @@ for epoch in range(epochs):
                             test_loss, test_acc])
 print("Done!")
 
-# saving trained model
+
 
 path = "/PATH/Contextualizer_mlp/Experiments_cifar10/weights_contextualizer"
 model_name = "ContextualizerImageClassification_cifar10"