Update train.py

train.py

@@ -15,23 +15,23 @@ Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
 
 training_data = datasets.CIFAR10(
-                                       root='data',
-                                       train=True,
-                                       download=True,
-                                       transform=transform 
+                                       root = 'data',
+                                       train = True,
+                                       download = True,
+                                       transform = transform 
                                        )
 
 test_data = datasets.CIFAR10(
-                                       root='data',
-                                       train=False,
-                                       download=True,
-                                       transform=transform 
+                                       root = 'data',
+                                       train = False,
+                                       download = True,
+                                       transform = transform 
                                        )                                       
                                       
 batch_size = 128
 
-train_dataloader = DataLoader(training_data, batch_size=batch_size,shuffle=True)
-test_dataloader = DataLoader(test_data, batch_size=batch_size)
+train_dataloader = DataLoader(training_data, batch_size = batch_size, shuffle = True)
+test_dataloader = DataLoader(test_data, batch_size = batch_size)
 
 for X, y in test_dataloader:
     print(f"Shape of X [N,C,H,W]:{X.shape}")
@@ -45,17 +45,17 @@ print(f"using {device} device")
 class TensorMapperImageClassification(TensorMapper):
     def __init__(
         self,
-        image_size=32,
-        patch_size=4,
-        in_channels=3,
-        num_classes=10,     
+        image_size = 32,
+        patch_size = 4,
+        in_channels = 3,
+        num_classes = 10,     
         d_model = 256,
-        num_layers=4,
+        num_layers = 4,
                
     ):
         super().__init__(d_model, num_layers)
         self.patcher = nn.Conv2d(
-            in_channels, d_model, kernel_size=patch_size, stride=patch_size
+            in_channels, d_model, kernel_size = patch_size, stride = patch_size
         )
         self.classifier = nn.Linear(d_model, num_classes)
 
@@ -66,7 +66,7 @@ class TensorMapperImageClassification(TensorMapper):
         patches = patches.permute(0, 2, 3, 1)
         patches = patches.view(batch_size, -1, num_channels)
         embedding = self.model(patches)
-        embedding = embedding.mean(dim=1) 
+        embedding = embedding.mean(dim = 1) 
         out = self.classifier(embedding)
         return out
 
@@ -74,7 +74,7 @@ model = TensorMapperImageClassification().to(device)
 print(model)
 
 loss_fn = nn.CrossEntropyLoss()
-optimizer = torch.optim.Adam(model.parameters(),lr=1e-3)
+optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
 
 def train(dataloader, model, loss_fn, optimizer):
     size = len(dataloader.dataset)
@@ -86,7 +86,7 @@ def train(dataloader, model, loss_fn, optimizer):
         X, y = X.to(device), y.to(device)
               
         pred = model(X)
-        loss = loss_fn(pred,y)
+        loss = loss_fn(pred, y)
                
         optimizer.zero_grad()
         loss.backward()
@@ -101,8 +101,8 @@ def train(dataloader, model, loss_fn, optimizer):
 
     train_loss /= num_batches
     train_accuracy = 100. * correct.item() / size
-    print(f"Train Metric: \n Accuracy: {(100*correct):>0.1f}% \n")
-    return train_loss,train_accuracy 
+    print(f"Train Metric: \n Accuracy: {(100 * correct):>0.1f}% \n")
+    return train_loss, train_accuracy 
 
 def test(dataloader, model, loss_fn):
     size = len(dataloader.dataset)            
@@ -118,8 +118,8 @@ def test(dataloader, model, loss_fn):
     	correct += (pred.argmax(1) == y).type(torch.float).sum().item()
     test_loss /= num_batches
     correct /= size
-    print(f"Test Metrics: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")  
-    test_accuracy = 100*correct      
+    print(f"Test Metrics: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")  
+    test_accuracy = 100 * correct      
     return test_loss, test_accuracy
 
 logname = "/PATH/Tensor_Mapper/Experiments_cifar10/logs_tensormapper/logs_cifar10.csv"
@@ -144,4 +144,4 @@ print("Done!")
 path = "/PATH/Tensor_Mapper/Experiments_cifar10/weights_tensormapper"
 model_name = "TensorMapperImageClassification_cifar10"
 torch.save(model.state_dict(), f"{path}/{model_name}.pth")
-print(f"Saved Model State to {path}/{model_name}.pth ")
+print(f"Saved Model State to {path}/{model_name}.pth ")