Adding presets in the classification reference scripts.

Author: datumbox

references/classification/README.md

@@ -124,22 +124,9 @@ Training converges at about 10 epochs.
 For post training quant, device is set to CPU. For training, the device is set to CUDA
 
 ### Command to evaluate quantized models using the pre-trained weights:
-For all quantized models except inception_v3:
+For all quantized models:
 ```
 python references/classification/train_quantization.py  --data-path='imagenet_full_size/' \
     --device='cpu' --test-only --backend='fbgemm' --model='<model_name>'
 ```
 
-For inception_v3, since it expects tensors with a size of N x 3 x 299 x 299, before running above command,
-need to change the input size of dataset_test in train.py to:
-```
-dataset_test = torchvision.datasets.ImageFolder(
-    valdir,
-    transforms.Compose([
-        transforms.Resize(342),
-        transforms.CenterCrop(299),
-        transforms.ToTensor(),
-        normalize,
-    ]))
-```
-

references/classification/presets.py

@@ -0,0 +1,38 @@
+from torchvision.transforms import autoaugment, transforms
+
+
+class ClassificationPresetTrain:
+    def __init__(self, crop_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), auto_augment_policy=None,
+                 random_erase_prob=0.0):
+        trans = [
+            transforms.RandomResizedCrop(crop_size),
+            transforms.RandomHorizontalFlip(),
+        ]
+        if auto_augment_policy is not None:
+            aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
+            trans.append(autoaugment.AutoAugment(policy=aa_policy))
+        trans.extend([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std),
+        ])
+        if random_erase_prob > 0:
+            trans.append(transforms.RandomErasing(p=random_erase_prob))
+
+        self.transforms = transforms.Compose(trans)
+
+    def __call__(self, img):
+        return self.transforms(img)
+
+
+class ClassificationPresetEval:
+    def __init__(self, crop_size, resize_size=256, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+
+        self.transforms = transforms.Compose([
+            transforms.Resize(resize_size),
+            transforms.CenterCrop(crop_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std),
+        ])
+
+    def __call__(self, img):
+        return self.transforms(img)

references/classification/train.py

@@ -6,8 +6,8 @@
 import torch.utils.data
 from torch import nn
 import torchvision
-from torchvision import transforms
 
+import presets
 import utils
 
 try:
@@ -82,8 +82,7 @@ def _get_cache_path(filepath):
 def load_data(traindir, valdir, args):
     # Data loading code
     print("Loading data")
-    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                                     std=[0.229, 0.224, 0.225])
+    resize_size, crop_size = (342, 299) if args.model == 'inception_v3' else (256, 224)
 
     print("Loading training data")
     st = time.time()
@@ -93,22 +92,10 @@ def load_data(traindir, valdir, args):
         print("Loading dataset_train from {}".format(cache_path))
         dataset, _ = torch.load(cache_path)
     else:
-        trans = [
-            transforms.RandomResizedCrop(224),
-            transforms.RandomHorizontalFlip(),
-        ]
-        if args.auto_augment is not None:
-            aa_policy = transforms.AutoAugmentPolicy(args.auto_augment)
-            trans.append(transforms.AutoAugment(policy=aa_policy))
-        trans.extend([
-            transforms.ToTensor(),
-            normalize,
-        ])
-        if args.random_erase > 0:
-            trans.append(transforms.RandomErasing(p=args.random_erase))
         dataset = torchvision.datasets.ImageFolder(
             traindir,
-            transforms.Compose(trans))
+            presets.ClassificationPresetTrain(crop_size=crop_size, auto_augment_policy=args.auto_augment,
+                                              random_erase_prob=args.random_erase))
         if args.cache_dataset:
             print("Saving dataset_train to {}".format(cache_path))
             utils.mkdir(os.path.dirname(cache_path))
@@ -124,12 +111,7 @@ def load_data(traindir, valdir, args):
     else:
         dataset_test = torchvision.datasets.ImageFolder(
             valdir,
-            transforms.Compose([
-                transforms.Resize(256),
-                transforms.CenterCrop(224),
-                transforms.ToTensor(),
-                normalize,
-            ]))
+            presets.ClassificationPresetEval(crop_size=crop_size, resize_size=resize_size))
         if args.cache_dataset:
             print("Saving dataset_test to {}".format(cache_path))
             utils.mkdir(os.path.dirname(cache_path))