Prototype references

references/classification/presets.py

@@ -1,5 +1,5 @@
 import torch
-from torchvision.transforms import autoaugment, transforms
+from torchvision.prototype import transforms
 from torchvision.transforms.functional import InterpolationMode
 
 
@@ -17,20 +17,20 @@ def __init__(
     ):
         trans = [transforms.RandomResizedCrop(crop_size, interpolation=interpolation)]
         if hflip_prob > 0:
-            trans.append(transforms.RandomHorizontalFlip(hflip_prob))
+            trans.append(transforms.RandomHorizontalFlip(p=hflip_prob))
         if auto_augment_policy is not None:
             if auto_augment_policy == "ra":
-                trans.append(autoaugment.RandAugment(interpolation=interpolation))
+                trans.append(transforms.RandAugment(interpolation=interpolation))
             elif auto_augment_policy == "ta_wide":
-                trans.append(autoaugment.TrivialAugmentWide(interpolation=interpolation))
+                trans.append(transforms.TrivialAugmentWide(interpolation=interpolation))
             elif auto_augment_policy == "augmix":
-                trans.append(autoaugment.AugMix(interpolation=interpolation))
+                trans.append(transforms.AugMix(interpolation=interpolation))
             else:
-                aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
-                trans.append(autoaugment.AutoAugment(policy=aa_policy, interpolation=interpolation))
+                aa_policy = transforms.AutoAugmentPolicy(auto_augment_policy)
+                trans.append(transforms.AutoAugment(policy=aa_policy, interpolation=interpolation))
         trans.extend(
             [
-                transforms.PILToTensor(),
+                transforms.ToImageTensor(),
                 transforms.ConvertImageDtype(torch.float),
                 transforms.Normalize(mean=mean, std=std),
             ]
@@ -59,7 +59,7 @@ def __init__(
             [
                 transforms.Resize(resize_size, interpolation=interpolation),
                 transforms.CenterCrop(crop_size),
-                transforms.PILToTensor(),
+                transforms.ToImageTensor(),
                 transforms.ConvertImageDtype(torch.float),
                 transforms.Normalize(mean=mean, std=std),
             ]

references/classification/train.py

@@ -4,14 +4,17 @@
 import warnings
 
 import presets
+from sampler import RASampler
+from transforms import WrapIntoFeatures
+import utils  # usort: skip
+
 import torch
 import torch.utils.data
 import torchvision
-import transforms
-import utils
-from sampler import RASampler
+
 from torch import nn
 from torch.utils.data.dataloader import default_collate
+from torchvision.prototype import features, transforms
 from torchvision.transforms.functional import InterpolationMode
 
 
@@ -128,12 +131,13 @@ def load_data(traindir, valdir, args):
         random_erase_prob = getattr(args, "random_erase", 0.0)
         dataset = torchvision.datasets.ImageFolder(
             traindir,
-            presets.ClassificationPresetTrain(
+            transform=presets.ClassificationPresetTrain(
                 crop_size=train_crop_size,
                 interpolation=interpolation,
                 auto_augment_policy=auto_augment_policy,
                 random_erase_prob=random_erase_prob,
             ),
+            target_transform=lambda target: features.Label(target),
         )
         if args.cache_dataset:
             print(f"Saving dataset_train to {cache_path}")
@@ -158,7 +162,8 @@ def load_data(traindir, valdir, args):
 
         dataset_test = torchvision.datasets.ImageFolder(
             valdir,
-            preprocessing,
+            transform=preprocessing,
+            target_transform=lambda target: features.Label(target),
         )
         if args.cache_dataset:
             print(f"Saving dataset_test to {cache_path}")
@@ -200,14 +205,21 @@ def main(args):
 
     collate_fn = None
     num_classes = len(dataset.classes)
-    mixup_transforms = []
+    mixup_or_cutmix = []
     if args.mixup_alpha > 0.0:
-        mixup_transforms.append(transforms.RandomMixup(num_classes, p=1.0, alpha=args.mixup_alpha))
+        mixup_or_cutmix.append(transforms.RandomMixup(alpha=args.mixup_alpha, p=1.0))
     if args.cutmix_alpha > 0.0:
-        mixup_transforms.append(transforms.RandomCutmix(num_classes, p=1.0, alpha=args.cutmix_alpha))
-    if mixup_transforms:
-        mixupcutmix = torchvision.transforms.RandomChoice(mixup_transforms)
-        collate_fn = lambda batch: mixupcutmix(*default_collate(batch))  # noqa: E731
+        mixup_or_cutmix.append(transforms.RandomCutmix(alpha=args.cutmix_alpha, p=1.0))
+    if mixup_or_cutmix:
+        batch_transform = transforms.Compose(
+            [
+                WrapIntoFeatures(),
+                transforms.LabelToOneHot(num_categories=num_classes),
+                transforms.ToDtype(torch.float, features.OneHotLabel),
+                transforms.RandomChoice(*mixup_or_cutmix),
+            ]
+        )
+        collate_fn = lambda batch: batch_transform(default_collate(batch))  # noqa: E731
     data_loader = torch.utils.data.DataLoader(
         dataset,
         batch_size=args.batch_size,

references/classification/transforms.py

@@ -1,183 +1,8 @@
-import math
-from typing import Tuple
+from torch import nn
+from torchvision.prototype import features
 
-import torch
-from torch import Tensor
-from torchvision.transforms import functional as F
 
-
-class RandomMixup(torch.nn.Module):
-    """Randomly apply Mixup to the provided batch and targets.
-    The class implements the data augmentations as described in the paper
-    `"mixup: Beyond Empirical Risk Minimization" <https://arxiv.org/abs/1710.09412>`_.
-
-    Args:
-        num_classes (int): number of classes used for one-hot encoding.
-        p (float): probability of the batch being transformed. Default value is 0.5.
-        alpha (float): hyperparameter of the Beta distribution used for mixup.
-            Default value is 1.0.
-        inplace (bool): boolean to make this transform inplace. Default set to False.
-    """
-
-    def __init__(self, num_classes: int, p: float = 0.5, alpha: float = 1.0, inplace: bool = False) -> None:
-        super().__init__()
-
-        if num_classes < 1:
-            raise ValueError(
-                f"Please provide a valid positive value for the num_classes. Got num_classes={num_classes}"
-            )
-
-        if alpha <= 0:
-            raise ValueError("Alpha param can't be zero.")
-
-        self.num_classes = num_classes
-        self.p = p
-        self.alpha = alpha
-        self.inplace = inplace
-
-    def forward(self, batch: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
-        """
-        Args:
-            batch (Tensor): Float tensor of size (B, C, H, W)
-            target (Tensor): Integer tensor of size (B, )
-
-        Returns:
-            Tensor: Randomly transformed batch.
-        """
-        if batch.ndim != 4:
-            raise ValueError(f"Batch ndim should be 4. Got {batch.ndim}")
-        if target.ndim != 1:
-            raise ValueError(f"Target ndim should be 1. Got {target.ndim}")
-        if not batch.is_floating_point():
-            raise TypeError(f"Batch dtype should be a float tensor. Got {batch.dtype}.")
-        if target.dtype != torch.int64:
-            raise TypeError(f"Target dtype should be torch.int64. Got {target.dtype}")
-
-        if not self.inplace:
-            batch = batch.clone()
-            target = target.clone()
-
-        if target.ndim == 1:
-            target = torch.nn.functional.one_hot(target, num_classes=self.num_classes).to(dtype=batch.dtype)
-
-        if torch.rand(1).item() >= self.p:
-            return batch, target
-
-        # It's faster to roll the batch by one instead of shuffling it to create image pairs
-        batch_rolled = batch.roll(1, 0)
-        target_rolled = target.roll(1, 0)
-
-        # Implemented as on mixup paper, page 3.
-        lambda_param = float(torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0])
-        batch_rolled.mul_(1.0 - lambda_param)
-        batch.mul_(lambda_param).add_(batch_rolled)
-
-        target_rolled.mul_(1.0 - lambda_param)
-        target.mul_(lambda_param).add_(target_rolled)
-
-        return batch, target
-
-    def __repr__(self) -> str:
-        s = (
-            f"{self.__class__.__name__}("
-            f"num_classes={self.num_classes}"
-            f", p={self.p}"
-            f", alpha={self.alpha}"
-            f", inplace={self.inplace}"
-            f")"
-        )
-        return s
-
-
-class RandomCutmix(torch.nn.Module):
-    """Randomly apply Cutmix to the provided batch and targets.
-    The class implements the data augmentations as described in the paper
-    `"CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features"
-    <https://arxiv.org/abs/1905.04899>`_.
-
-    Args:
-        num_classes (int): number of classes used for one-hot encoding.
-        p (float): probability of the batch being transformed. Default value is 0.5.
-        alpha (float): hyperparameter of the Beta distribution used for cutmix.
-            Default value is 1.0.
-        inplace (bool): boolean to make this transform inplace. Default set to False.
-    """
-
-    def __init__(self, num_classes: int, p: float = 0.5, alpha: float = 1.0, inplace: bool = False) -> None:
-        super().__init__()
-        if num_classes < 1:
-            raise ValueError("Please provide a valid positive value for the num_classes.")
-        if alpha <= 0:
-            raise ValueError("Alpha param can't be zero.")
-
-        self.num_classes = num_classes
-        self.p = p
-        self.alpha = alpha
-        self.inplace = inplace
-
-    def forward(self, batch: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
-        """
-        Args:
-            batch (Tensor): Float tensor of size (B, C, H, W)
-            target (Tensor): Integer tensor of size (B, )
-
-        Returns:
-            Tensor: Randomly transformed batch.
-        """
-        if batch.ndim != 4:
-            raise ValueError(f"Batch ndim should be 4. Got {batch.ndim}")
-        if target.ndim != 1:
-            raise ValueError(f"Target ndim should be 1. Got {target.ndim}")
-        if not batch.is_floating_point():
-            raise TypeError(f"Batch dtype should be a float tensor. Got {batch.dtype}.")
-        if target.dtype != torch.int64:
-            raise TypeError(f"Target dtype should be torch.int64. Got {target.dtype}")
-
-        if not self.inplace:
-            batch = batch.clone()
-            target = target.clone()
-
-        if target.ndim == 1:
-            target = torch.nn.functional.one_hot(target, num_classes=self.num_classes).to(dtype=batch.dtype)
-
-        if torch.rand(1).item() >= self.p:
-            return batch, target
-
-        # It's faster to roll the batch by one instead of shuffling it to create image pairs
-        batch_rolled = batch.roll(1, 0)
-        target_rolled = target.roll(1, 0)
-
-        # Implemented as on cutmix paper, page 12 (with minor corrections on typos).
-        lambda_param = float(torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0])
-        _, H, W = F.get_dimensions(batch)
-
-        r_x = torch.randint(W, (1,))
-        r_y = torch.randint(H, (1,))
-
-        r = 0.5 * math.sqrt(1.0 - lambda_param)
-        r_w_half = int(r * W)
-        r_h_half = int(r * H)
-
-        x1 = int(torch.clamp(r_x - r_w_half, min=0))
-        y1 = int(torch.clamp(r_y - r_h_half, min=0))
-        x2 = int(torch.clamp(r_x + r_w_half, max=W))
-        y2 = int(torch.clamp(r_y + r_h_half, max=H))
-
-        batch[:, :, y1:y2, x1:x2] = batch_rolled[:, :, y1:y2, x1:x2]
-        lambda_param = float(1.0 - (x2 - x1) * (y2 - y1) / (W * H))
-
-        target_rolled.mul_(1.0 - lambda_param)
-        target.mul_(lambda_param).add_(target_rolled)
-
-        return batch, target
-
-    def __repr__(self) -> str:
-        s = (
-            f"{self.__class__.__name__}("
-            f"num_classes={self.num_classes}"
-            f", p={self.p}"
-            f", alpha={self.alpha}"
-            f", inplace={self.inplace}"
-            f")"
-        )
-        return s
+class WrapIntoFeatures(nn.Module):
+    def forward(self, sample):
+        input, target = sample
+        return features.Image(input), features.Label(target)