change effinet, transformations and add num_workers to dataloader

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}}/models/EffiNet.py

@@ -19,14 +19,12 @@
 class EffiNet(ArtModule):
     def __init__(self, num_classes=100, lr=1e-3):
         super().__init__()
-        self.model = timm.create_model('tf_efficientnet_b4.in1k', pretrained=True)
+        self.model = timm.create_model('efficientnet_b2.ra_in1kx    ', pretrained=True, num_classes=100)
         self.loss = torch.nn.CrossEntropyLoss()
         self.lr = lr
-        self.model.classifier = nn.Linear(self.model.classifier.in_features, num_classes)
         self.preprocess = transforms.Compose([
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
             transforms.Resize(256),
-            transforms.CenterCrop(224),
         ])
 
     def parse_data(self, data):