Merge branch 'main' into ra-reps

Author: datumbox

references/classification/README.md

@@ -125,7 +125,7 @@ torchrun --nproc_per_node=8 train.py\
 ```
 Here `$MODEL` is one of `regnet_x_400mf`, `regnet_x_800mf`, `regnet_x_1_6gf`, `regnet_y_400mf`, `regnet_y_800mf` and `regnet_y_1_6gf`. Please note we used learning rate 0.4 for `regent_y_400mf` to get the same Acc@1 as [the paper)(https://arxiv.org/abs/2003.13678).
 
-### Medium models
+#### Medium models
 ```
 torchrun --nproc_per_node=8 train.py\
      --model $MODEL --epochs 100 --batch-size 64 --wd 0.00005 --lr=0.4\
@@ -134,7 +134,7 @@ torchrun --nproc_per_node=8 train.py\
 ```
 Here `$MODEL` is one of `regnet_x_3_2gf`, `regnet_x_8gf`, `regnet_x_16gf`, `regnet_y_3_2gf` and `regnet_y_8gf`.
 
-### Large models
+#### Large models
 ```
 torchrun --nproc_per_node=8 train.py\
      --model $MODEL --epochs 100 --batch-size 32 --wd 0.00005 --lr=0.2\
@@ -143,6 +143,28 @@ torchrun --nproc_per_node=8 train.py\
 ```
 Here `$MODEL` is one of `regnet_x_32gf`, `regnet_y_16gf` and `regnet_y_32gf`.
 
+### Vision Transformer
+
+#### Base models
+```
+torchrun --nproc_per_node=8 train.py\
+    --model $MODEL --epochs 300 --batch-size 64 --opt adamw --lr 0.003 --wd 0.3\
+    --lr-scheduler cosineannealinglr --lr-warmup-method linear --lr-warmup-epochs 30\
+    --lr-warmup-decay 0.033 --amp --label-smoothing 0.11 --mixup-alpha 0.2 --auto-augment ra\
+    --clip-grad-norm 1 --ra-sampler --cutmix-alpha 1.0 --model-ema
+```
+Here `$MODEL` is one of `vit_b_16` and `vit_b_32`.
+
+#### Large models
+```
+torchrun --nproc_per_node=8 train.py\
+    --model $MODEL --epochs 300 --batch-size 16 --opt adamw --lr 0.003 --wd 0.3\
+    --lr-scheduler cosineannealinglr --lr-warmup-method linear --lr-warmup-epochs 30\
+    --lr-warmup-decay 0.033 --amp --label-smoothing 0.11 --mixup-alpha 0.2 --auto-augment ra\
+    --clip-grad-norm 1 --ra-sampler --cutmix-alpha 1.0 --model-ema
+```
+Here `$MODEL` is one of `vit_l_16` and `vit_l_32`.
+
 ## Mixed precision training
 Automatic Mixed Precision (AMP) training on GPU for Pytorch can be enabled with the [torch.cuda.amp](https://pytorch.org/docs/stable/amp.html?highlight=amp#module-torch.cuda.amp).
 

references/optical_flow/README.md

@@ -10,7 +10,14 @@ training and evaluation scripts to quickly bootstrap research.
 The RAFT large model was trained on Flying Chairs and then on Flying Things.
 Both used 8 A100 GPUs and a batch size of 2 (so effective batch size is 16). The
 rest of the hyper-parameters are exactly the same as the original RAFT training
-recipe from https://github.com/princeton-vl/RAFT.
+recipe from https://github.com/princeton-vl/RAFT. The original recipe trains for
+100000 updates (or steps) on each dataset - this corresponds to about 72 and 20
+epochs on Chairs and Things respectively:
+
+```
+num_epochs = ceil(num_steps / number_of_steps_per_epoch)
+           = ceil(num_steps / (num_samples / effective_batch_size))
+```
 
 ```
 torchrun --nproc_per_node 8 --nnodes 1 train.py \
@@ -21,7 +28,7 @@ torchrun --nproc_per_node 8 --nnodes 1 train.py \
     --batch-size 2 \
     --lr 0.0004 \
     --weight-decay 0.0001 \
-    --num-steps 100000 \
+    --epochs 72 \
     --output-dir $chairs_dir
 ```
 
@@ -34,7 +41,7 @@ torchrun --nproc_per_node 8 --nnodes 1 train.py \
     --batch-size 2 \
     --lr 0.000125 \
     --weight-decay 0.0001 \
-    --num-steps 100000 \
+    --epochs 20 \
     --freeze-batch-norm \
     --output-dir $things_dir\
     --resume $chairs_dir/$name_chairs.pth

references/optical_flow/train.py

@@ -1,5 +1,6 @@
 import argparse
 import warnings
+from math import ceil
 from pathlib import Path
 
 import torch
@@ -168,7 +169,7 @@ def validate(model, args):
             warnings.warn(f"Can't validate on {val_dataset}, skipping.")
 
 
-def train_one_epoch(model, optimizer, scheduler, train_loader, logger, current_step, args):
+def train_one_epoch(model, optimizer, scheduler, train_loader, logger, args):
     for data_blob in logger.log_every(train_loader):
 
         optimizer.zero_grad()
@@ -189,13 +190,6 @@ def train_one_epoch(model, optimizer, scheduler, train_loader, logger, current_s
         optimizer.step()
         scheduler.step()
 
-        current_step += 1
-
-        if current_step == args.num_steps:
-            return True, current_step
-
-    return False, current_step
-
 
 def main(args):
     utils.setup_ddp(args)
@@ -243,7 +237,8 @@ def main(args):
     scheduler = torch.optim.lr_scheduler.OneCycleLR(
         optimizer=optimizer,
         max_lr=args.lr,
-        total_steps=args.num_steps + 100,
+        epochs=args.epochs,
+        steps_per_epoch=ceil(len(train_dataset) / (args.world_size * args.batch_size)),
         pct_start=0.05,
         cycle_momentum=False,
         anneal_strategy="linear",
@@ -252,26 +247,22 @@ def main(args):
     logger = utils.MetricLogger()
 
     done = False
-    current_epoch = current_step = 0
-    while not done:
+    for current_epoch in range(args.epochs):
         print(f"EPOCH {current_epoch}")
 
         sampler.set_epoch(current_epoch)  # needed, otherwise the data loading order would be the same for all epochs
-        done, current_step = train_one_epoch(
+        train_one_epoch(
             model=model,
             optimizer=optimizer,
             scheduler=scheduler,
             train_loader=train_loader,
             logger=logger,
-            current_step=current_step,
             args=args,
         )
 
         # Note: we don't sync the SmoothedValues across processes, so the printed metrics are just those of rank 0
         print(f"Epoch {current_epoch} done. ", logger)
 
-        current_epoch += 1
-
         if args.rank == 0:
             # TODO: Also save the optimizer and scheduler
             torch.save(model.state_dict(), Path(args.output_dir) / f"{args.name}_{current_epoch}.pth")
@@ -310,10 +301,8 @@ def get_args_parser(add_help=True):
     )
     parser.add_argument("--val-dataset", type=str, nargs="+", help="The dataset(s) to use for validation.")
     parser.add_argument("--val-freq", type=int, default=2, help="Validate every X epochs")
-    # TODO: eventually, it might be preferable to support epochs instead of num_steps.
-    # Keeping it this way for now to reproduce results more easily.
-    parser.add_argument("--num-steps", type=int, default=100000, help="The total number of steps (updates) to train.")
-    parser.add_argument("--batch-size", type=int, default=6)
+    parser.add_argument("--epochs", type=int, default=20, help="The total number of epochs to train.")
+    parser.add_argument("--batch-size", type=int, default=2)
 
     parser.add_argument("--lr", type=float, default=0.00002, help="Learning rate for AdamW optimizer")
     parser.add_argument("--weight-decay", type=float, default=0.00005, help="Weight decay for AdamW optimizer")

test/test_prototype_builtin_datasets.py

@@ -2,9 +2,11 @@
 
 import builtin_dataset_mocks
 import pytest
+import torch
+from torch.utils.data.datapipes.iter.grouping import ShardingFilterIterDataPipe as ShardingFilter
 from torch.utils.data.graph import traverse
-from torchdata.datapipes.iter import IterDataPipe
-from torchvision.prototype import datasets, features
+from torchdata.datapipes.iter import IterDataPipe, Shuffler
+from torchvision.prototype import datasets, transforms
 from torchvision.prototype.datasets._api import DEFAULT_DECODER
 from torchvision.prototype.utils._internal import sequence_to_str
 
@@ -88,15 +90,36 @@ def test_decoding(self, dataset, mock_info):
             )
 
     @dataset_parametrization(decoder=DEFAULT_DECODER)
-    def test_at_least_one_feature(self, dataset, mock_info):
-        sample = next(iter(dataset))
-        if not any(isinstance(value, features.Feature) for value in sample.values()):
-            raise AssertionError("The sample contained no feature.")
+    def test_no_vanilla_tensors(self, dataset, mock_info):
+        vanilla_tensors = {key for key, value in next(iter(dataset)).items() if type(value) is torch.Tensor}
+        if vanilla_tensors:
+            raise AssertionError(
+                f"The values of key(s) "
+                f"{sequence_to_str(sorted(vanilla_tensors), separate_last='and ')} contained vanilla tensors."
+            )
+
+    @dataset_parametrization()
+    def test_transformable(self, dataset, mock_info):
+        next(iter(dataset.map(transforms.Identity())))
 
     @dataset_parametrization()
     def test_traversable(self, dataset, mock_info):
         traverse(dataset)
 
+    @dataset_parametrization()
+    @pytest.mark.parametrize("annotation_dp_type", (Shuffler, ShardingFilter), ids=lambda type: type.__name__)
+    def test_has_annotations(self, dataset, mock_info, annotation_dp_type):
+        def scan(graph):
+            for node, sub_graph in graph.items():
+                yield node
+                yield from scan(sub_graph)
+
+        for dp in scan(traverse(dataset)):
+            if type(dp) is annotation_dp_type:
+                break
+        else:
+            raise AssertionError(f"The dataset doesn't comprise a {annotation_dp_type.__name__}() datapipe.")
+
 
 class TestQMNIST:
     @pytest.mark.parametrize(

torchvision/csrc/io/video/video.cpp

@@ -188,7 +188,7 @@ Video::Video(std::string videoPath, std::string stream, int64_t numThreads) {
   c10::Dict<std::string, std::vector<double>> ccMetadata;
   c10::Dict<std::string, std::vector<double>> subsMetadata;
 
-  // calback and metadata defined in struct
+  // callback and metadata defined in struct
   succeeded = decoder.init(params, std::move(callback), &metadata);
   if (succeeded) {
     for (const auto& header : metadata) {
@@ -254,7 +254,7 @@ bool Video::setCurrentStream(std::string stream = "video") {
       numThreads_ // global number of threads
   );
 
-  // calback and metadata defined in Video.h
+  // callback and metadata defined in Video.h
   return (decoder.init(params, std::move(callback), &metadata));
 }
 
@@ -280,7 +280,7 @@ void Video::Seek(double ts, bool fastSeek = false) {
       numThreads_ // global number of threads
   );
 
-  // calback and metadata defined in Video.h
+  // callback and metadata defined in Video.h
   succeeded = decoder.init(params, std::move(callback), &metadata);
   LOG(INFO) << "Decoder init at seek " << succeeded << "\n";
 }

torchvision/csrc/ops/deform_conv2d.cpp

@@ -20,6 +20,7 @@ at::Tensor deform_conv2d(
     int64_t groups,
     int64_t offset_groups,
     bool use_mask) {
+  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.deform_conv2d.deform_conv2d");
   static auto op = c10::Dispatcher::singleton()
                        .findSchemaOrThrow("torchvision::deform_conv2d", "")
                        .typed<decltype(deform_conv2d)>();

torchvision/csrc/ops/nms.cpp

@@ -9,6 +9,7 @@ at::Tensor nms(
     const at::Tensor& dets,
     const at::Tensor& scores,
     double iou_threshold) {
+  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.nms.nms");
   static auto op = c10::Dispatcher::singleton()
                        .findSchemaOrThrow("torchvision::nms", "")
                        .typed<decltype(nms)>();

torchvision/csrc/ops/ps_roi_align.cpp

@@ -12,6 +12,7 @@ std::tuple<at::Tensor, at::Tensor> ps_roi_align(
     int64_t pooled_height,
     int64_t pooled_width,
     int64_t sampling_ratio) {
+  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.ps_roi_align.ps_roi_align");
   static auto op = c10::Dispatcher::singleton()
                        .findSchemaOrThrow("torchvision::ps_roi_align", "")
                        .typed<decltype(ps_roi_align)>();

torchvision/csrc/ops/ps_roi_pool.cpp

@@ -11,6 +11,7 @@ std::tuple<at::Tensor, at::Tensor> ps_roi_pool(
     double spatial_scale,
     int64_t pooled_height,
     int64_t pooled_width) {
+  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.ps_roi_pool.ps_roi_pool");
   static auto op = c10::Dispatcher::singleton()
                        .findSchemaOrThrow("torchvision::ps_roi_pool", "")
                        .typed<decltype(ps_roi_pool)>();

torchvision/csrc/ops/roi_align.cpp

@@ -16,6 +16,7 @@ at::Tensor roi_align(
     bool aligned) // The flag for pixel shift
 // along each axis.
 {
+  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.roi_align.roi_align");
   static auto op = c10::Dispatcher::singleton()
                        .findSchemaOrThrow("torchvision::roi_align", "")
                        .typed<decltype(roi_align)>();

torchvision/csrc/ops/roi_pool.cpp

@@ -11,6 +11,7 @@ std::tuple<at::Tensor, at::Tensor> roi_pool(
     double spatial_scale,
     int64_t pooled_height,
     int64_t pooled_width) {
+  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.roi_pool.roi_pool");
   static auto op = c10::Dispatcher::singleton()
                        .findSchemaOrThrow("torchvision::roi_pool", "")
                        .typed<decltype(roi_pool)>();

torchvision/datasets/vision.py

@@ -35,7 +35,7 @@ def __init__(
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,
     ) -> None:
-        _log_api_usage_once("datasets", self.__class__.__name__)
+        _log_api_usage_once(self)
         if isinstance(root, torch._six.string_classes):
             root = os.path.expanduser(root)
         self.root = root

torchvision/extension.py

@@ -1,3 +1,8 @@
+import ctypes
+import os
+import sys
+from warnings import warn
+
 import torch
 
 from ._internally_replaced_utils import _get_extension_path
@@ -67,4 +72,22 @@ def _check_cuda_version():
     return _version
 
 
+def _load_library(lib_name):
+    lib_path = _get_extension_path(lib_name)
+    # On Windows Python-3.8+ has `os.add_dll_directory` call,
+    # which is called from _get_extension_path to configure dll search path
+    # Condition below adds a workaround for older versions by
+    # explicitly calling `LoadLibraryExW` with the following flags:
+    #  - LOAD_LIBRARY_SEARCH_DEFAULT_DIRS (0x1000)
+    #  - LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR (0x100)
+    if os.name == "nt" and sys.version_info < (3, 8):
+        _kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+        if hasattr(_kernel32, "LoadLibraryExW"):
+            _kernel32.LoadLibraryExW(lib_path, None, 0x00001100)
+        else:
+            warn("LoadLibraryExW is missing in kernel32.dll")
+
+    torch.ops.load_library(lib_path)
+
+
 _check_cuda_version()

torchvision/io/_video_opt.py

@@ -5,12 +5,11 @@
 
 import torch
 
-from .._internally_replaced_utils import _get_extension_path
+from ..extension import _load_library
 
 
 try:
-    lib_path = _get_extension_path("video_reader")
-    torch.ops.load_library(lib_path)
+    _load_library("video_reader")
     _HAS_VIDEO_OPT = True
 except (ImportError, OSError):
     _HAS_VIDEO_OPT = False

torchvision/io/image.py

@@ -1,15 +1,15 @@
 from enum import Enum
+from warnings import warn
 
 import torch
 
-from .._internally_replaced_utils import _get_extension_path
+from ..extension import _load_library
 
 
 try:
-    lib_path = _get_extension_path("image")
-    torch.ops.load_library(lib_path)
-except (ImportError, OSError):
-    pass
+    _load_library("image")
+except (ImportError, OSError) as e:
+    warn(f"Failed to load image Python extension: {e}")
 
 
 class ImageReadMode(Enum):

torchvision/models/alexnet.py

@@ -18,7 +18,7 @@
 class AlexNet(nn.Module):
     def __init__(self, num_classes: int = 1000, dropout: float = 0.5) -> None:
         super().__init__()
-        _log_api_usage_once("models", self.__class__.__name__)
+        _log_api_usage_once(self)
         self.features = nn.Sequential(
             nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
             nn.ReLU(inplace=True),

torchvision/models/densenet.py

@@ -163,7 +163,7 @@ def __init__(
     ) -> None:
 
         super().__init__()
-        _log_api_usage_once("models", self.__class__.__name__)
+        _log_api_usage_once(self)
 
         # First convolution
         self.features = nn.Sequential(

torchvision/models/detection/generalized_rcnn.py

@@ -27,7 +27,7 @@ class GeneralizedRCNN(nn.Module):
 
     def __init__(self, backbone: nn.Module, rpn: nn.Module, roi_heads: nn.Module, transform: nn.Module) -> None:
         super().__init__()
-        _log_api_usage_once("models", self.__class__.__name__)
+        _log_api_usage_once(self)
         self.transform = transform
         self.backbone = backbone
         self.rpn = rpn