1st version of tutorial + new model

Author: mmaecki

{{cookiecutter.project_slug}}/MyDataset.py

@@ -3,24 +3,25 @@
 from torch.utils.data import DataLoader
 
 class CifarDataModule(pl.LightningDataModule):
-    def __init__(self, batch_size: int = 32):
+    def __init__(self, batch_size: int = 32, num_workers: int = 4):
         super().__init__()
         self.batch_size = batch_size
         self.dataset = load_dataset("cifar100").with_format("torch")
         self.dataset = self.dataset.rename_columns({"img": "input", "fine_label": "target"})
         # self.dataset = self.dataset.rename_columns({"img": "input", "coarse_label": "target"})
         self.dataset = self.dataset.remove_columns(["coarse_label"])
         # self.dataset = self.dataset.remove_columns(["fine_label"])
+        self.num_workers = num_workers
 
     def setup(self, stage: str):
         self.train = self.dataset["train"]
         self.test = self.dataset["test"]
 
     def train_dataloader(self):
-        return DataLoader(self.dataset["train"], batch_size=self.batch_size)
+        return DataLoader(self.dataset["train"], batch_size=self.batch_size, num_workers=self.num_workers)
 
     def val_dataloader(self):
-        return DataLoader(self.dataset["test"], batch_size=self.batch_size)
+        return DataLoader(self.dataset["test"], batch_size=self.batch_size, num_workers=self.num_workers)
 
     def log_params(self):
         return {

{{cookiecutter.project_slug}}/Tutorial.ipynb

@@ -0,0 +1,59 @@
+from typing import Dict
+from art.core import ArtModule
+import torch
+import timm
+import torch.nn as nn
+from torchvision import transforms
+import numpy as np
+from einops import rearrange
+from art.utils.enums import (
+    BATCH,
+    INPUT,
+    LOSS,
+    PREDICTION,
+    TARGET,
+    TRAIN_LOSS,
+    VALIDATION_LOSS,
+)
+
+class EffiNet(ArtModule):
+    def __init__(self, num_classes=100, lr=1e-3):
+        super().__init__()
+        self.model = timm.create_model('efficientnet_b2.ra_in1k', pretrained=True, num_classes=100)
+        self.loss = torch.nn.CrossEntropyLoss()
+        self.lr = lr
+        self.preprocess = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            transforms.Resize(256),
+        ])
+
+    def parse_data(self, data):
+        """This is first step of your pipeline it always has batch keys inside"""
+        X = data[BATCH][INPUT]
+        X = X / 255
+        X = rearrange(X, "b h w c -> b c h w")
+        X = self.preprocess(X)
+        target = data[BATCH][TARGET].long()
+        return {INPUT: X, TARGET: target}
+    
+
+    
+    def predict(self, data: Dict):       
+        return {PREDICTION: self.model(data[INPUT]), TARGET: data[TARGET]}
+    
+    def compute_loss(self, data):
+        # Notice that the loss calculation is done in MetricsCalculator!
+        # We only need to specify which loss (metric) we want to use
+        loss = data["CrossEntropyLoss"]
+        return {LOSS: loss}
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+    
+    def log_params(self):
+        # Log relevant parameters
+        return {
+            "lr": self.lr,
+            "model_name": self.model.__class__.__name__,
+            "n_parameters": sum(p.numel() for p in self.parameters() if p.requires_grad),
+        }

{{cookiecutter.project_slug}}/models/EffiNet.py

@@ -15,7 +15,7 @@
     VALIDATION_LOSS,
 )
 
-class ResNet18(ArtModule):
+class ResNet(ArtModule):
     def __init__(self, num_classes=100, lr=1e-3):
         super().__init__()
         self.model = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models', 'resnet18_swsl')
@@ -27,8 +27,6 @@ def __init__(self, num_classes=100, lr=1e-3):
             transforms.Resize(256),
             transforms.CenterCrop(224),
         ])
-        # for name, para in self.model.named_parameters():
-        #     para.requires_grad = True
 
     def parse_data(self, data):
         """This is first step of your pipeline it always has batch keys inside"""

{{cookiecutter.project_slug}}/models/ResNet.py

@@ -20,16 +20,11 @@ def do(self, previous_states):
         # Now tell me what are the names of these classes
         class_names = list(self.datamodule.dataset["train"].features[TARGET].names)
 
-        class_counts = Counter(targets)
-
         # Now calculate number of images in each class
-        number_of_examples_in_each_class = [
-            class_counts[i] for i in range(number_of_classes)
-        ]
+        class_counts = Counter(targets)
 
         # Now tell me dimensions of each image
         img_dimensions = self.datamodule.train_dataloader().dataset[0][INPUT].shape
-        figures = []
         for cls in class_names:
             class_indices = [i for i, label in enumerate(targets) if label == cls]
             class_samples = np.random.choice(class_indices, 5, replace=False).tolist()
@@ -47,15 +42,13 @@ def do(self, previous_states):
             MatplotLibSaver().save(
                 fig, self.get_full_step_name(), self.get_class_image_path(cls)
             )
-            figures.append(fig)
 
         self.results.update(
             {
                 "number_of_classes": number_of_classes,
                 "class_names": class_names,
-                "number_of_examples_in_each_class": number_of_examples_in_each_class,
+                "number_of_examples_in_each_class": class_counts,
                 "img_dimensions": img_dimensions,
-                "images": figures,
             }
         )
     def log_params(self):