loop

Author: mrkito

config.py

@@ -1,16 +1,15 @@
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
 
-batch_size=64
+batch_size = 64
 
-train_dir='/home/deepkot/Downloads/imagenette2-160/train/'
-val_dir='/home/deepkot/Downloads/imagenette2-160/val/'
+train_dir = 'imagenette2-160/train/'
+val_dir = 'imagenette2-160/val/'
 
 
 def get_train_transforms(size):
     TRAIN_TRANSFORMS = A.Compose([
         A.Resize(size, size),
-
         A.HorizontalFlip(p=0.5),
         A.Rotate(limit=30),
         A.ImageCompression(),
@@ -46,7 +45,6 @@ def get_train_transforms(size):
 
 def get_val_transforms(size):
     VALID_TRANSFORMS = A.Compose([
-
         A.Resize(size, size),
         A.Normalize(),
         ToTensorV2(),

train_loop.py

@@ -1,48 +1,16 @@
-import os
-import torch
 from catalyst.contrib.nn import RAdam, Lookahead
 from torch.nn import functional as F
-from torch.utils.data import DataLoader
 from catalyst import dl
-from catalyst.contrib.data.transforms import ToTensor
-from catalyst.contrib.datasets import MNIST
 from catalyst.utils import metrics, set_global_seed, prepare_cudnn
-
-
-from catalyst.utils import (
-    create_dataset, create_dataframe, get_dataset_labeling, map_dataframe
-)
-# train_dir='/home/deepkot/Downloads/imagenette2-160/train/n01440764'
-# dataset = create_dataset(dirs=f"{train_dir}/*", extension="*.jpg")
-# df = create_dataframe(dataset, columns=["class", "filepath"])
-#
-# tag_to_label = get_dataset_labeling(df, "class")
-# class_names = [
-#     name for name, id_ in sorted(tag_to_label.items(), key=lambda x: x[1])
-# ]
-#
-# df_with_labels = map_dataframe(
-#     df,
-#     tag_column="class",
-#     class_column="label",
-#     tag2class=tag_to_label,
-#     verbose=False
-# )
-# df_with_labels.head()
 from config import train_dir, val_dir, batch_size, get_train_transforms, get_val_transforms
 from dataset_bn import get_loaders
 from model_batchnorm import resnet18
 
 
-def train_research_bn(model):
-    SEED = 69
-    set_global_seed(SEED)
-    prepare_cudnn(deterministic=True)
-
+def train_research_bn(model, log):
     optimizer = Lookahead(RAdam(model.parameters(), lr=0.02))
 
-
-    train_data_transforms, val_data_transforms = get_train_transforms(224), get_val_transforms(224)
+    train_data_transforms, val_data_transforms = get_train_transforms(112), get_val_transforms(112)
 
     loaders = get_loaders(
         train_dir=train_dir,
@@ -52,17 +20,14 @@ def train_research_bn(model):
         batch_size=batch_size,
     )
 
-
-
     class CustomRunner(dl.Runner):
 
         def predict_batch(self, batch):
             # model inference step
             return self.model(batch[0].to(self.device).view(batch[0].size(0), -1))
 
         def _handle_batch(self, batch):
-            # model train/valid step
-            x, y = batch['image'],batch['label'],
+            x, y = batch['image'], batch['label'],
             y_hat = self.model(x)
 
             loss = F.cross_entropy(y_hat, y)
@@ -77,18 +42,31 @@ def _handle_batch(self, batch):
                 self.state.optimizer.zero_grad()
 
     runner = CustomRunner()
-    # model training
     runner.train(
         model=model,
         optimizer=optimizer,
         loaders=loaders,
-        logdir="./logs",
+        logdir=log,
         num_epochs=15,
         verbose=True,
         load_best_on_end=True,
     )
-    return  runner
+    return runner
 
 
 if __name__ == '__main__':
-    train_research_bn(resnet18(train_bn=False))
+    SEED = 69
+    set_global_seed(SEED)
+    prepare_cudnn(deterministic=True)
+
+    model = resnet18(train_bn=True)
+    train_research_bn(model, "resnet")
+
+    model = resnet18(train_bn=False)
+    train_research_bn(model, 'resnet_no_bn')
+
+    model = resnet18(train_bn=True)
+    for name, p in model.named_parameters():
+        if 'bn' not in name:
+            p.requires_grad = False
+    train_research_bn(model, 'resnet_bn_only')