Update S10.py

S10.py

@@ -9,7 +9,7 @@ os.environ['KMP_DUPLICATE_LIB_OK']='True'
 import pytorch_lightening as pl
 
 #MODEL
-class LtAutoEncoder(L.LighteningModule):
+class LitAutoEncoder(L.LighteningModule):
     def __init__(self):
         super(Net,self).__init__()
         self.conv1 = nn.Sequential (
@@ -86,72 +86,73 @@ class LtAutoEncoder(L.LighteningModule):
         x = x.view(-1, 10)
         return x
 
-    def training_step(self, batch, batch_idx):
-        # training_step defines the train loop.
-        # It is independent of forward
-        train_losses = []
-        train_acc = []
-        model.train()
-        pbar = tqdm(train_loader)
-        correct = 0
-        processed = 0
-        criterion = nn.CrossEntropyLoss()
-        for batch_idx, (data, target) in enumerate(pbar):
-        # get samples
-            data, target = data.to(device), target.to(device)
-            loss = criterion(y_pred, target)
-       
-            train_losses.append(loss)
-            optimizer.zero_grad()
-            # Backpropagation
-            loss.backward()
-            optimizer.step()
-
-            pred = y_pred.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
-            correct += pred.eq(target.view_as(pred)).sum().item()
-            processed += len(data)
-
-            pbar.set_description(desc= f'Loss={loss.item()} Batch_id={batch_idx} Accuracy={100*correct/processed:0.2f}')
-            train_acc.append(100*correct/processed)
-        
-        
-    def testing_step(self, batch, batch_idx):
-        test_losses = []
-        test_acc = []
-        model.eval()
-        test_loss = 0
-        correct = 0
-        with torch.no_grad():
-            for data, target in test_loader:
-                data, target = data.to(device), target.to(device)
-                output = model(data)
-                pred =output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
-                correct += pred.eq(target.view_as(pred)).sum().item()
-
-        test_loss /= len(test_loader.dataset)
-        test_losses.append(test_loss)
-
-        print('\nTest set:  Accuracy: {}/{} ({:.2f}%)\n'.format(
-             correct, len(test_loader.dataset),
-            100. * correct / len(test_loader.dataset)))
-
-        test_acc.append(100. * correct / len(test_loader.dataset))   
-        
+    
+    def cross_entropy_loss(self, logits, labels):
+    return F.nll_loss(logits, labels)
+
+      def training_step(self, train_batch, batch_idx):
+          x, y = train_batch
+          logits = self.forward(x)
+          loss = self.cross_entropy_loss(logits, y)
+          self.log('train_loss', loss)
+          return loss
+
+
+      def validation_step(self, val_batch, batch_idx):
+          x, y = val_batch
+          logits = self.forward(x)
+          loss = self.cross_entropy_loss(logits, y)
+          self.log('val_loss', loss)
         
     #OPTIMIZER
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
         return optimizer
 
-model=LtAutoEncoder(Net)  
+ 
 
 
-# setup data
-dataset = CIFAR10(os.getcwd(), download=True, transform=ToTensor())
-train_loader = utils.data.DataLoader(dataset)
+class DataModule(pl.LightningDataModule):
 
-# train the model (hint: here are some helpful Trainer arguments for rapid idea iteration)
-trainer = pl.Trainer(limit_train_batches=100, max_epochs=1)
-trainer.fit(model=autoencoder, train_dataloaders=train_loader)
-    
+  def setup(self, stage):
+    # transforms for images
+    train_transform=A.Compose([
+        A.Normalize(
+            mean = (0.4914, 0.4822, 0.4465),
+            std = (0.2470, 0.2435, 0.2616), always_apply = True
+        ),
+        A.HorizontalFlip(),
+        A.ShiftScaleRotate(shift_limit=(-0.2, 0.2), scale_limit=(-0.2, 0.2), rotate_limit=(-15, 15), p=0.5),
+        A.PadIfNeeded(min_height=36, min_width=36, p=1.0),
+        A.RandomCrop (32, 32, always_apply=False, p=1.0),
+        A.CenterCrop(32, 32, always_apply=False, p=1.0),
+        A.CoarseDropout(max_holes = 1, max_height=8, max_width=8, min_holes = 1, min_height=8, min_width=8,
+                        fill_value=(0.4914, 0.4822, 0.4465), always_apply=False, p=0.5),
+
+        ToTensorV2()
+        ])
+
+    valid_transform=A.Compose([
+         A.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010), always_apply = True),
+                                  ToTensorV2()
+        ])
+
+      
+    # prepare transforms standard to MNIST
+    self.mnist_train = CIFAR10(os.getcwd(), train=True, download=True, transform=train_transform)
+    self.mnist_test = CIFAR10(os.getcwd(), train=False, download=True, transform=valid_transform)
+
+  def train_dataloader(self):
+    return DataLoader(self.mnist_train, batch_size=512)
+
+  def val_dataloader(self):
+    return DataLoader(self.mnist_test, batch_size=512)
+
+data_module = DataModule()
+
+# train
+model=LitAutoEncoder() 
+trainer = pl.Trainer()
+
+trainer.fit(model, data_module)