Update S10.py

S10.py

@@ -2,16 +2,67 @@
 import torch
 import torch.nn as nn
 import pytorch_lightning as pl
+from torchvision.datasets import MNIST
 from torchvision.transforms import ToTensor
+from torch.utils.data import DataLoader, random_split
 import torch
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
-
 from torchvision import transforms
 import numpy as np
+import torch
 from torchvision import datasets
 from torch.utils.data import Dataset, DataLoader
 from torchvision.transforms import ToTensor
+from torchmetrics import Accuracy
+import matplotlib.pyplot as plt
+from torch.nn import functional as F
+
+
+class DataAugmentationTrain():
+    """Module to perform data augmentation using Kornia on torch tensors."""
+
+
+    def __init__(self):
+      self.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(p=1.0),
+          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()
+          ])
+
+    def __call__(self, image):
+      image = np.array(image)
+      image = self.train_transform(image = image)['image']
+      return image
+
+class DataAugmentationTest():
+    """Module to perform data augmentation using Kornia on torch tensors."""
+
+
+    def __init__(self):
+      self.test_transform = A.Compose([
+          A.Normalize(
+              mean = (0.4914, 0.4822, 0.4465),
+              std = (0.2470, 0.2435, 0.2616), always_apply = True
+          ),
+          ToTensorV2()
+          ])
+
+    def __call__(self, image):
+      image = np.array(image)
+      img = self.test_transform(image = image)['image']
+      return image
 
 
 torch.manual_seed(1)
@@ -71,6 +122,7 @@ class MyModel(nn.Module):
         self.fc = nn.Linear(512, 10, bias=True)
 
     def forward(self, x):
+        #x = x.unsqueeze(0)
         x = self.conv1(x)
 
         x = self.conv11(x)
@@ -89,148 +141,94 @@ class MyModel(nn.Module):
         #x = x.randn(512, 1)
 
 # squeeze the tensor to size 512x
-        x = x.squeeze(dim=[2, 3])
+        #x = x.squeeze(dim=[2, 3])
+        x = x.squeeze(dim=2)
+        x = x.squeeze(dim=2)
+        #x = x.squeeze(dim=2).squeeze(dim=3)
+        #x = x.squeeze(dim=2)
 
         #x = x.view(512, 10)
 
         x = self.fc(x)
 
         x = x.view(-1, 10)
+        x = F.log_softmax(x, dim=-1)
         return x
 
-class MyLightningModule(pl.LightningModule):
-    def __init__(self):
+class Cifar10SearchDataset(pl.LightningModule):
+    def __init__(self, train_transform=None, test_transform=None):
         super().__init__()
+        # not the best model: expereiment yourself
         self.model = MyModel()
+        #self.preprocess = Preprocess()  # per sample transforms
+        self.train_transform = DataAugmentationTrain()
+        self.test_transform = DataAugmentationTest()
+          # per batch augmentation_kornia
+        self.train_accuracy = Accuracy(task="multiclass", num_classes=10)
+        self.val_accuracy = Accuracy(task="multiclass", num_classes=10)
+
+    def forward(self, x):
+        return self.model(x)
 
     def training_step(self, batch, batch_idx):
-        x, y = batch
-        y_hat = self.model(x)
-        loss = nn.functional.cross_entropy(y_hat, y)
+        x, y= batch
+        logits = self.model(x)
+        loss = F.nll_loss(logits, y)
+        preds = torch.argmax(logits, dim=1)
+        self.train_accuracy(preds, y)
+        self.log("train_loss", loss, prog_bar=True)
+        self.log("train_acc", self.train_accuracy, prog_bar=True)
         return loss
 
     def validation_step(self, batch, batch_idx):
-        x, y = batch
-        y_hat = self.model(x)
-        loss = nn.functional.cross_entropy(y_hat, y)
-        acc = (y_hat.argmax(dim=1) == y).float().mean()
-        self.log('val_loss', loss)
-        self.log('val_acc', acc)
+        x, y= batch
+       
+        logits = self.model(x)
+        loss = F.nll_loss(logits, y)
+        preds = torch.argmax(logits, dim=1)
+        self.val_accuracy(preds, y)
+        self.log("val_loss", loss, prog_bar=True)
+        self.log("val_acc", self.val_accuracy, prog_bar=True)
+        #self.log('val_loss', loss)
+        #self.log('val_acc', acc)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        # Here we just reuse the validation_step for testing
+        return self.validation_step(batch, batch_idx)
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
         return optimizer
 
-def get_a_train_transform():
-    """Get transformer for training data
-
-    Returns:
-        Compose: Composed transformations
-    """
-    return 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()
-        ])
-
-
-
-def get_a_test_transform():
-    """Get transformer for test data
-
-    Returns:
-        Compose: Composed transformations
-    """
-    return A.Compose([
-         A.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010), always_apply = True),
-                                  ToTensorV2()
-    ])
-
-class Cifar10SearchDataset(Dataset):
-    """
-    Custom Dataset Class
-
-    """
-
-    def __init__(self, dataset, transforms=None):
-        """Initialize Dataset
-
-        Args:
-            dataset (Dataset): Pytorch Dataset instance
-            transforms (Transform.Compose, optional): Tranform function instance. Defaults to None.
-        """
-        self.transforms = transforms
-        self.dataset = dataset
-
-    def __len__(self):
-        """Get dataset length
-
-        Returns:
-            int: Length of dataset
-        """
-        return len(self.dataset)
-
-    def __getitem__(self, idx):
-        """Get an item form dataset
-
-        Args:
-            idx (int): id of item in dataset
-
-        Returns:
-            (tensor, int): Return tensor of transformer image, label
-        """
-        # Read Image and Label
-        image, label = self.dataset[idx]#,self.targets[idx]
-
-        image = np.array(image)
-
-        # Apply Transforms
-        if self.transforms is not None:
-            image = (self.transforms(image=image))["image"]
-
-        
-        return (image, label)
-
-
-def get_loader(train_transform, test_transform, batch_size=128, use_cuda=True):
-    """Get instance of tran and test loaders
-    self.batch_size = 128
-        self.device = device
-        self.use_cuda = use_cuda
-    Args:
-        train_transform (Transform): Instance of transform function for training
-        test_transform (Transform): Instance of transform function for validation
-        batch_size (int, optional): batch size to be uised in training. Defaults to 64.
-        use_cuda (bool, optional): Enable/Disable Cuda Gpu. Defaults to True.
-
-    Returns:
-        (DataLoader, DataLoader): Get instance of train and test data loaders
-    """
-    kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
-
-    train_loader = DataLoader(
-        Cifar10SearchDataset(datasets.CIFAR10('../data', train=True,
-                     download=True), transforms=train_transform),
-        batch_size=batch_size, shuffle=True, **kwargs)
-
-    test_loader = DataLoader(
-        Cifar10SearchDataset(datasets.CIFAR10('../data', train=False,
-                     download=True), transforms=test_transform),
-        batch_size=batch_size, shuffle=True, **kwargs)
-
-    return train_loader, test_loader
-
-use_cuda = torch.cuda.is_available()
-device = torch.device("cuda" if use_cuda else "cpu") 
-train_loader, test_loader = get_loader(get_a_train_transform(), get_a_test_transform(), batch_size=512, use_cuda=use_cuda)
+    def prepare_data(self):
+        # download data if needed
+        datasets.CIFAR10(root='./data', train=True, download=True)
+        datasets.CIFAR10(root='./data', train=False, download=True)
+        pass
+
+    def setup(self, stage=None):
+        # split data into train and test sets
+        if stage == "fit" or stage is None:
+          cifar10_full = datasets.CIFAR10(root='./data', train=True, download=True, transform = self.train_transform)
+          self.cifar10_train, self.cifar10_val = random_split(cifar10_full, [45000, 5000])
+
+        if stage == "test" or stage is None:
+          self.cifar10_test = datasets.CIFAR10(root='./data', train=False, download=True, transform = self.test_transform)
+
+    def train_dataloader(self):
+        return DataLoader(self.cifar10_train, batch_size=64)
+
+    def val_dataloader(self):
+        return DataLoader(self.cifar10_val, batch_size=64)
+
+    def test_dataloader(self):
+        return DataLoader(self.cifar10_test, batch_size=64)
+
+
+model = Cifar10SearchDataset()
+trainer = pl.Trainer(accelerator="auto",max_epochs=1)
+#datasetcifar=Cifar10SearchDataset()
+#trainer = pl.Trainer()
+trainer.fit(model)