Refactoring to use Compose.

Author: datumbox

references/classification/presets.py

@@ -1,43 +1,37 @@
-from torch import nn, Tensor
-from torchvision.transforms import autoaugment, transforms
+from torchvision.transforms import *
 
 
-class ClassificationPresetTrain(nn.Module):
+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):
-        super().__init__()
-
         trans = [
-            transforms.RandomResizedCrop(crop_size),
-            transforms.RandomHorizontalFlip(),
+            RandomResizedCrop(crop_size),
+            RandomHorizontalFlip(),
         ]
         if auto_augment_policy is not None:
-            aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
-            trans.append(autoaugment.AutoAugment(policy=aa_policy))
+            aa_policy = AutoAugmentPolicy(auto_augment_policy)
+            trans.append(AutoAugment(policy=aa_policy))
         trans.extend([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=mean, std=std),
+            ToTensor(),
+            Normalize(mean=mean, std=std),
         ])
         if random_erase_prob > 0:
-            trans.append(transforms.RandomErasing(p=random_erase_prob))
+            trans.append(RandomErasing(p=random_erase_prob))
 
-        self.transforms = nn.Sequential(*trans)
+        self.transforms = Compose(trans)
 
-    def forward(self, img: Tensor) -> Tensor:
+    def __call__(self, img):
         return self.transforms(img)
 
 
-class ClassificationPresetEval(nn.Module):
+class ClassificationPresetEval:
     def __init__(self, crop_size, resize_size=256, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
-        super().__init__()
-
-        trans = [
-            transforms.Resize(resize_size),
-            transforms.CenterCrop(crop_size),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=mean, std=std),
-        ]
-        self.transforms = nn.Sequential(*trans)
+        self.transforms = Compose([
+            Resize(resize_size),
+            CenterCrop(crop_size),
+            ToTensor(),
+            Normalize(mean=mean, std=std),
+        ])
 
-    def forward(self, img: Tensor) -> Tensor:
+    def __call__(self, img):
         return self.transforms(img)