[fbsync] Refactor preset transforms (#5562)

Summary:
* Refactor preset transforms

* Making presets public.

Reviewed By: vmoens

Differential Revision: D34878982

fbshipit-source-id: f1551f5d1f98a58a6820d28b4bd8a0dcbdc53226

Author: datumbox

references/classification/train.py

@@ -163,7 +163,7 @@ def load_data(traindir, valdir, args):
                 weights = prototype.models.get_weight(args.weights)
                 preprocessing = weights.transforms()
             else:
-                preprocessing = prototype.transforms.ImageNetEval(
+                preprocessing = prototype.transforms.ImageClassificationEval(
                     crop_size=val_crop_size, resize_size=val_resize_size, interpolation=interpolation
                 )
 

references/detection/train.py

@@ -57,7 +57,7 @@ def get_transform(train, args):
             weights = prototype.models.get_weight(args.weights)
             return weights.transforms()
         else:
-            return prototype.transforms.CocoEval()
+            return prototype.transforms.ObjectDetectionEval()
 
 
 def get_args_parser(add_help=True):

references/optical_flow/train.py

@@ -137,7 +137,7 @@ def validate(model, args):
             weights = prototype.models.get_weight(args.weights)
             preprocessing = weights.transforms()
         else:
-            preprocessing = prototype.transforms.RaftEval()
+            preprocessing = prototype.transforms.OpticalFlowEval()
     else:
         preprocessing = OpticalFlowPresetEval()
 

references/segmentation/train.py

@@ -42,7 +42,7 @@ def get_transform(train, args):
             weights = prototype.models.get_weight(args.weights)
             return weights.transforms()
         else:
-            return prototype.transforms.VocEval(resize_size=520)
+            return prototype.transforms.SemanticSegmentationEval(resize_size=520)
 
 
 def criterion(inputs, target):

references/video_classification/train.py

@@ -157,7 +157,7 @@ def main(args):
             weights = prototype.models.get_weight(args.weights)
             transform_test = weights.transforms()
         else:
-            transform_test = prototype.transforms.Kinect400Eval(crop_size=(112, 112), resize_size=(128, 171))
+            transform_test = prototype.transforms.VideoClassificationEval(crop_size=(112, 112), resize_size=(128, 171))
 
     if args.cache_dataset and os.path.exists(cache_path):
         print(f"Loading dataset_test from {cache_path}")

torchvision/prototype/models/alexnet.py

@@ -1,7 +1,7 @@
 from functools import partial
 from typing import Any, Optional
 
-from torchvision.prototype.transforms import ImageNetEval
+from torchvision.prototype.transforms import ImageClassificationEval
 from torchvision.transforms.functional import InterpolationMode
 
 from ...models.alexnet import AlexNet
@@ -16,7 +16,7 @@
 class AlexNet_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/alexnet-owt-7be5be79.pth",
-        transforms=partial(ImageNetEval, crop_size=224),
+        transforms=partial(ImageClassificationEval, crop_size=224),
         meta={
             "task": "image_classification",
             "architecture": "AlexNet",

torchvision/prototype/models/convnext.py

@@ -1,7 +1,7 @@
 from functools import partial
 from typing import Any, List, Optional
 
-from torchvision.prototype.transforms import ImageNetEval
+from torchvision.prototype.transforms import ImageClassificationEval
 from torchvision.transforms.functional import InterpolationMode
 
 from ...models.convnext import ConvNeXt, CNBlockConfig
@@ -56,7 +56,7 @@ def _convnext(
 class ConvNeXt_Tiny_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/convnext_tiny-983f1562.pth",
-        transforms=partial(ImageNetEval, crop_size=224, resize_size=236),
+        transforms=partial(ImageClassificationEval, crop_size=224, resize_size=236),
         meta={
             **_COMMON_META,
             "num_params": 28589128,
@@ -70,7 +70,7 @@ class ConvNeXt_Tiny_Weights(WeightsEnum):
 class ConvNeXt_Small_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/convnext_small-0c510722.pth",
-        transforms=partial(ImageNetEval, crop_size=224, resize_size=230),
+        transforms=partial(ImageClassificationEval, crop_size=224, resize_size=230),
         meta={
             **_COMMON_META,
             "num_params": 50223688,
@@ -84,7 +84,7 @@ class ConvNeXt_Small_Weights(WeightsEnum):
 class ConvNeXt_Base_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/convnext_base-6075fbad.pth",
-        transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
+        transforms=partial(ImageClassificationEval, crop_size=224, resize_size=232),
         meta={
             **_COMMON_META,
             "num_params": 88591464,
@@ -98,7 +98,7 @@ class ConvNeXt_Base_Weights(WeightsEnum):
 class ConvNeXt_Large_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/convnext_large-ea097f82.pth",
-        transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
+        transforms=partial(ImageClassificationEval, crop_size=224, resize_size=232),
         meta={
             **_COMMON_META,
             "num_params": 197767336,

torchvision/prototype/models/densenet.py

@@ -3,7 +3,7 @@
 from typing import Any, Optional, Tuple
 
 import torch.nn as nn
-from torchvision.prototype.transforms import ImageNetEval
+from torchvision.prototype.transforms import ImageClassificationEval
 from torchvision.transforms.functional import InterpolationMode
 
 from ...models.densenet import DenseNet
@@ -78,7 +78,7 @@ def _densenet(
 class DenseNet121_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/densenet121-a639ec97.pth",
-        transforms=partial(ImageNetEval, crop_size=224),
+        transforms=partial(ImageClassificationEval, crop_size=224),
         meta={
             **_COMMON_META,
             "num_params": 7978856,
@@ -92,7 +92,7 @@ class DenseNet121_Weights(WeightsEnum):
 class DenseNet161_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/densenet161-8d451a50.pth",
-        transforms=partial(ImageNetEval, crop_size=224),
+        transforms=partial(ImageClassificationEval, crop_size=224),
         meta={
             **_COMMON_META,
             "num_params": 28681000,
@@ -106,7 +106,7 @@ class DenseNet161_Weights(WeightsEnum):
 class DenseNet169_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/densenet169-b2777c0a.pth",
-        transforms=partial(ImageNetEval, crop_size=224),
+        transforms=partial(ImageClassificationEval, crop_size=224),
         meta={
             **_COMMON_META,
             "num_params": 14149480,
@@ -120,7 +120,7 @@ class DenseNet169_Weights(WeightsEnum):
 class DenseNet201_Weights(WeightsEnum):
     IMAGENET1K_V1 = Weights(
         url="https://download.pytorch.org/models/densenet201-c1103571.pth",
-        transforms=partial(ImageNetEval, crop_size=224),
+        transforms=partial(ImageClassificationEval, crop_size=224),
         meta={
             **_COMMON_META,
             "num_params": 20013928,

torchvision/prototype/models/detection/faster_rcnn.py

@@ -1,7 +1,7 @@
 from typing import Any, Optional, Union
 
 from torch import nn
-from torchvision.prototype.transforms import CocoEval
+from torchvision.prototype.transforms import ObjectDetectionEval
 from torchvision.transforms.functional import InterpolationMode
 
 from ....models.detection.faster_rcnn import (
@@ -43,7 +43,7 @@
 class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum):
     COCO_V1 = Weights(
         url="https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth",
-        transforms=CocoEval,
+        transforms=ObjectDetectionEval,
         meta={
             **_COMMON_META,
             "num_params": 41755286,
@@ -57,7 +57,7 @@ class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum):
 class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
     COCO_V1 = Weights(
         url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_fpn-fb6a3cc7.pth",
-        transforms=CocoEval,
+        transforms=ObjectDetectionEval,
         meta={
             **_COMMON_META,
             "num_params": 19386354,
@@ -71,7 +71,7 @@ class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
 class FasterRCNN_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
     COCO_V1 = Weights(
         url="https://download.pytorch.org/models/fasterrcnn_mobilenet_v3_large_320_fpn-907ea3f9.pth",
-        transforms=CocoEval,
+        transforms=ObjectDetectionEval,
         meta={
             **_COMMON_META,
             "num_params": 19386354,

torchvision/prototype/models/detection/fcos.py

@@ -1,7 +1,7 @@
 from typing import Any, Optional
 
 from torch import nn
-from torchvision.prototype.transforms import CocoEval
+from torchvision.prototype.transforms import ObjectDetectionEval
 from torchvision.transforms.functional import InterpolationMode
 
 from ....models.detection.fcos import (
@@ -27,7 +27,7 @@
 class FCOS_ResNet50_FPN_Weights(WeightsEnum):
     COCO_V1 = Weights(
         url="https://download.pytorch.org/models/fcos_resnet50_fpn_coco-99b0c9b7.pth",
-        transforms=CocoEval,
+        transforms=ObjectDetectionEval,
         meta={
             "task": "image_object_detection",
             "architecture": "FCOS",