Update train.py

train.py

@@ -7,8 +7,6 @@ from torchvision import datasets
 from torchvision.transforms import ToTensor, Normalize, Compose 
 from tensormapper import TensorMapper
 
-
-
 transform = Compose([
  
 ToTensor(),
@@ -29,28 +27,21 @@ test_data = datasets.CIFAR10(
                                        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)
 
-
 for X, y in test_dataloader:
     print(f"Shape of X [N,C,H,W]:{X.shape}")
     print(f"Shape of y:{y.shape}{y.dtype}")
     break
 
-
-
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print(f"using {device} device") 
 
-
-
 class TensorMapperImageClassification(TensorMapper):
     def __init__(
         self,
@@ -60,8 +51,7 @@ class TensorMapperImageClassification(TensorMapper):
         num_classes=10,     
         d_model = 256,
         num_layers=4,
-       
-        
+               
     ):
         super().__init__(d_model, num_layers)
         self.patcher = nn.Conv2d(
@@ -83,14 +73,9 @@ class TensorMapperImageClassification(TensorMapper):
 model = TensorMapperImageClassification().to(device)
 print(model)
 
-
-
 loss_fn = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters(),lr=1e-3)
 
-
-
-
 def train(dataloader, model, loss_fn, optimizer):
     size = len(dataloader.dataset)
     num_batches = len(dataloader)
@@ -99,12 +84,10 @@ def train(dataloader, model, loss_fn, optimizer):
     correct = 0
     for batch, (X,y) in enumerate(dataloader):
         X, y = X.to(device), y.to(device)
-       
-       
+              
         pred = model(X)
         loss = loss_fn(pred,y)
-        
-       
+               
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
@@ -112,9 +95,6 @@ def train(dataloader, model, loss_fn, optimizer):
         _, labels = torch.max(pred.data, 1)
         correct += labels.eq(y.data).type(torch.float).sum()
 
-        
-
-
         if batch % 100 == 0:
             loss, current = loss.item(), batch * len(X)
             print(f"loss: {loss:>7f}   [{current:>5d}/{size:>5d}]")
@@ -124,10 +104,6 @@ def train(dataloader, model, loss_fn, optimizer):
     print(train_accuracy)
     return train_loss,train_accuracy 
 
-
-
-
-
 def test(dataloader, model, loss_fn):
     size = len(dataloader.dataset)            
     num_batches = len(dataloader)
@@ -146,10 +122,6 @@ def test(dataloader, model, loss_fn):
     test_accuracy = 100*correct      
     return test_loss, test_accuracy
 
-
-
-
-
 logname = "/PATH/Tensor_Mapper/Experiments_cifar10/logs_tensormapper/logs_cifar10.csv"
 if not os.path.exists(logname):
   with open(logname, 'w') as logfile:
@@ -157,7 +129,6 @@ if not os.path.exists(logname):
     logwriter.writerow(['epoch', 'train loss', 'train acc',
                         'test loss', 'test acc'])
 
-
 epochs = 100
 for epoch in range(epochs):
     print(f"Epoch {epoch+1}\n-----------------------------------")
@@ -167,12 +138,10 @@ for epoch in range(epochs):
         logwriter = csv.writer(logfile, delimiter=',')
         logwriter.writerow([epoch+1, train_loss, train_acc,
                             test_loss, test_acc])
+        
 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 ")
-