Merge branch 'main' into checkpoint

Author: yiwen-song

.circleci/config.yml

@@ -311,7 +311,7 @@ jobs:
           descr: Install Python type check utilities
       - run:
           name: Check Python types statically
-          command: mypy --config-file mypy.ini
+          command: mypy --install-types --non-interactive --config-file mypy.ini
 
   unittest_torchhub:
     docker:

.circleci/config.yml.in

@@ -32,8 +32,8 @@ def build_workflows(prefix="", filter_branch=None, upload=False, indentation=6,
         for os_type in ["linux", "macos", "win"]:
             python_versions = PYTHON_VERSIONS
             cu_versions_dict = {
-                "linux": ["cpu", "cu102", "cu111", "cu113", "rocm4.2", "rocm4.3.1"],
-                "win": ["cpu", "cu111", "cu113"],
+                "linux": ["cpu", "cu102", "cu111", "cu113", "cu115", "rocm4.2", "rocm4.3.1"],
+                "win": ["cpu", "cu111", "cu113", "cu115"],
                 "macos": ["cpu"],
             }
             cu_versions = cu_versions_dict[os_type]
@@ -128,6 +128,7 @@ def upload_doc_job(filter_branch):
     "cu111": "pytorch/manylinux-cuda111",
     "cu112": "pytorch/manylinux-cuda112",
     "cu113": "pytorch/manylinux-cuda113",
+    "cu115": "pytorch/manylinux-cuda115",
 }
 
 

.circleci/regenerate.py

@@ -11,5 +11,11 @@ setup_conda_pytorch_constraint
 setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
 setup_junit_results_folder
-# nvidia channel included for cudatoolkit >= 11
-conda build -c defaults -c nvidia $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision
+
+# nvidia channel included for cudatoolkit >= 11 however for 11.5 we use conda-forge
+export CUDATOOLKIT_CHANNEL="nvidia"
+if [[ "$CU_VERSION" == cu115 ]]; then
+    export CUDATOOLKIT_CHANNEL="conda-forge"
+fi
+
+conda build -c defaults -c $CUDATOOLKIT_CHANNEL $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

packaging/build_conda.sh

@@ -46,6 +46,14 @@ setup_cuda() {
 
   # Now work out the CUDA settings
   case "$CU_VERSION" in
+    cu115)
+      if [[ "$OSTYPE" == "msys" ]]; then
+        export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.5"
+      else
+        export CUDA_HOME=/usr/local/cuda-11.5/
+      fi
+      export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
+      ;;
     cu113)
       if [[ "$OSTYPE" == "msys" ]]; then
         export CUDA_HOME="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.3"
@@ -289,6 +297,9 @@ setup_conda_cudatoolkit_constraint() {
     export CONDA_BUILD_VARIANT="cpu"
   else
     case "$CU_VERSION" in
+      cu115)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.5,<11.6 # [not osx]"
+        ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.3,<11.4 # [not osx]"
         ;;
@@ -333,6 +344,9 @@ setup_conda_cudatoolkit_plain_constraint() {
     export CMAKE_USE_CUDA=0
   else
     case "$CU_VERSION" in
+      cu115)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.5"
+        ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.3"
         ;;

packaging/pkg_helpers.bash

@@ -24,6 +24,7 @@ requirements:
   run:
     - python
     - defaults::numpy >=1.11
+    - requests
     - libpng
     - ffmpeg >=4.2  # [not win]
     - jpeg

packaging/torchvision/meta.yaml

@@ -21,6 +21,7 @@ set CUDA_VER_MAJOR=%CUDA_VER:~0,-1%
 set CUDA_VER_MINOR=%CUDA_VER:~-1,1%
 set CUDA_VERSION_STR=%CUDA_VER_MAJOR%.%CUDA_VER_MINOR%
 
+
 if %CUDA_VER% EQU 92 goto cuda92
 if %CUDA_VER% EQU 100 goto cuda100
 if %CUDA_VER% EQU 101 goto cuda101
@@ -29,6 +30,8 @@ if %CUDA_VER% EQU 110 goto cuda110
 if %CUDA_VER% EQU 111 goto cuda111
 if %CUDA_VER% EQU 112 goto cuda112
 if %CUDA_VER% EQU 113 goto cuda113
+if %CUDA_VER% EQU 115 goto cuda115
+
 
 echo CUDA %CUDA_VERSION_STR% is not supported
 exit /b 1
@@ -180,6 +183,25 @@ if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
 
 goto cuda_common
 
+:cuda115
+
+set CUDA_INSTALL_EXE=cuda_11.5.0_496.13_win10.exe
+if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
+    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    if errorlevel 1 exit /b 1
+    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    set "ARGS=thrust_11.5 nvcc_11.5 cuobjdump_11.5 nvprune_11.5 nvprof_11.5 cupti_11.5 cublas_11.5 cublas_dev_11.5 cudart_11.5 cufft_11.5 cufft_dev_11.5 curand_11.5 curand_dev_11.5 cusolver_11.5 cusolver_dev_11.5 cusparse_11.5 cusparse_dev_11.5 npp_11.5 npp_dev_11.5 nvrtc_11.5 nvrtc_dev_11.5 nvml_dev_11.5"
+)
+
+set CUDNN_INSTALL_ZIP=cudnn-11.3-windows-x64-v8.2.0.53.zip
+if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
+    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+    if errorlevel 1 exit /b 1
+    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+)
+
+goto cuda_common
+
 :cuda_common
 
 if not exist "%SRC_DIR%\temp_build\NvToolsExt.7z" (
@@ -222,7 +244,7 @@ set "NVTOOLSEXT_PATH=%ProgramFiles%\NVIDIA Corporation\NvToolsExt\bin\x64"
 if not exist "%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\nvcc.exe" (
     echo CUDA %CUDA_VERSION_STR% installed failed.
     echo --------- RunDll32.exe.log
-    type "%SRC_DIR%\temp_build\cuda\cuda_install_logs\LOG.RunDll32.exe.log"    
+    type "%SRC_DIR%\temp_build\cuda\cuda_install_logs\LOG.RunDll32.exe.log"
     echo --------- setup.exe.log -------
     type "%SRC_DIR%\temp_build\cuda\cuda_install_logs\LOG.setup.exe.log"
     exit /b 1

packaging/windows/internal/cuda_install.bat

@@ -15,7 +15,7 @@ class RASampler(torch.utils.data.Sampler):
     https://github.com/facebookresearch/deit/blob/main/samplers.py
     """
 
-    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
+    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True, seed=0):
         if num_replicas is None:
             if not dist.is_available():
                 raise RuntimeError("Requires distributed package to be available!")
@@ -32,11 +32,12 @@ def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
         self.total_size = self.num_samples * self.num_replicas
         self.num_selected_samples = int(math.floor(len(self.dataset) // 256 * 256 / self.num_replicas))
         self.shuffle = shuffle
+        self.seed = seed
 
     def __iter__(self):
         # Deterministically shuffle based on epoch
         g = torch.Generator()
-        g.manual_seed(self.epoch)
+        g.manual_seed(self.seed + self.epoch)
         if self.shuffle:
             indices = torch.randperm(len(self.dataset), generator=g).tolist()
         else:

references/classification/sampler.py

@@ -9,7 +9,7 @@
 import torchvision
 import transforms
 import utils
-from references.classification.sampler import RASampler
+from sampler import RASampler
 from torch import nn
 from torch.utils.data.dataloader import default_collate
 from torchvision.transforms.functional import InterpolationMode

references/classification/train.py

@@ -0,0 +1,65 @@
+# Optical flow reference training scripts
+
+This folder contains reference training scripts for optical flow.
+They serve as a log of how to train specific models, so as to provide baseline
+training and evaluation scripts to quickly bootstrap research.
+
+
+### RAFT Large
+
+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.
+
+```
+torchrun --nproc_per_node 8 --nnodes 1 train.py \
+    --dataset-root $dataset_root \
+    --name $name_chairs \
+    --model raft_large \
+    --train-dataset chairs \
+    --batch-size 2 \
+    --lr 0.0004 \
+    --weight-decay 0.0001 \
+    --num-steps 100000 \
+    --output-dir $chairs_dir
+```
+
+```
+torchrun --nproc_per_node 8 --nnodes 1 train.py \
+    --dataset-root $dataset_root \
+    --name $name_things \
+    --model raft_large \
+    --train-dataset things \
+    --batch-size 2 \
+    --lr 0.000125 \
+    --weight-decay 0.0001 \
+    --num-steps 100000 \
+    --freeze-batch-norm \
+    --output-dir $things_dir\
+    --resume $chairs_dir/$name_chairs.pth
+```
+
+
+### Evaluation
+
+```
+torchrun --nproc_per_node 1 --nnodes 1 train.py --val-dataset sintel --batch-size 1 --dataset-root $dataset_root --model raft_large --pretrained
+```
+
+This should give an epe of about 1.3822 on the clean pass and 2.7161 on the
+final pass of Sintel-train. Results may vary slightly depending on the batch
+size and the number of GPUs. For the most accurate resuts use 1 GPU and
+`--batch-size 1`:
+
+```
+Sintel val clean epe: 1.3822	1px: 0.9028	3px: 0.9573	5px: 0.9697	per_image_epe: 1.3822	f1: 4.0248
+Sintel val final epe: 2.7161	1px: 0.8528	3px: 0.9204	5px: 0.9392	per_image_epe: 2.7161	f1: 7.5964
+```
+
+You can also evaluate on Kitti train:
+
+```
+torchrun --nproc_per_node 1 --nnodes 1 train.py --val-dataset kitti --batch-size 1 --dataset-root $dataset_root --model raft_large --pretrained
+Kitti val epe: 4.7968	1px: 0.6388	3px: 0.8197	5px: 0.8661	per_image_epe: 4.5118	f1: 16.0679
+```

references/optical_flow/README.md

@@ -3,10 +3,16 @@
 from pathlib import Path
 
 import torch
+import torchvision.models.optical_flow
 import utils
 from presets import OpticalFlowPresetTrain, OpticalFlowPresetEval
 from torchvision.datasets import KittiFlow, FlyingChairs, FlyingThings3D, Sintel, HD1K
-from torchvision.models.optical_flow import raft_large, raft_small
+
+try:
+    from torchvision.prototype import models as PM
+    from torchvision.prototype.models import optical_flow as PMOF
+except ImportError:
+    PM = PMOF = None
 
 
 def get_train_dataset(stage, dataset_root):
@@ -125,6 +131,13 @@ def inner_loop(blob):
 
 def validate(model, args):
     val_datasets = args.val_dataset or []
+
+    if args.weights:
+        weights = PM.get_weight(args.weights)
+        preprocessing = weights.transforms()
+    else:
+        preprocessing = OpticalFlowPresetEval()
+
     for name in val_datasets:
         if name == "kitti":
             # Kitti has different image sizes so we need to individually pad them, we can't batch.
@@ -134,14 +147,14 @@ def validate(model, args):
                     f"Batch-size={args.batch_size} was passed. For technical reasons, evaluating on Kitti can only be done with a batch-size of 1."
                 )
 
-            val_dataset = KittiFlow(root=args.dataset_root, split="train", transforms=OpticalFlowPresetEval())
+            val_dataset = KittiFlow(root=args.dataset_root, split="train", transforms=preprocessing)
             _validate(
                 model, args, val_dataset, num_flow_updates=24, padder_mode="kitti", header="Kitti val", batch_size=1
             )
         elif name == "sintel":
             for pass_name in ("clean", "final"):
                 val_dataset = Sintel(
-                    root=args.dataset_root, split="train", pass_name=pass_name, transforms=OpticalFlowPresetEval()
+                    root=args.dataset_root, split="train", pass_name=pass_name, transforms=preprocessing
                 )
                 _validate(
                     model,
@@ -187,7 +200,11 @@ def train_one_epoch(model, optimizer, scheduler, train_loader, logger, current_s
 def main(args):
     utils.setup_ddp(args)
 
-    model = raft_small() if args.small else raft_large()
+    if args.weights:
+        model = PMOF.__dict__[args.model](weights=args.weights)
+    else:
+        model = torchvision.models.optical_flow.__dict__[args.model](pretrained=args.pretrained)
+
     model = model.to(args.local_rank)
     model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
 
@@ -306,7 +323,12 @@ def get_args_parser(add_help=True):
         "--freeze-batch-norm", action="store_true", help="Set BatchNorm modules of the model in eval mode."
     )
 
-    parser.add_argument("--small", action="store_true", help="Use the 'small' RAFT architecture.")
+    parser.add_argument(
+        "--model", type=str, default="raft_large", help="The name of the model to use - either raft_large or raft_small"
+    )
+    # TODO: resume, pretrained, and weights should be in an exclusive arg group
+    parser.add_argument("--pretrained", action="store_true", help="Whether to use pretrained weights")
+    parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load.")
 
     parser.add_argument(
         "--num_flow_updates",

references/optical_flow/train.py

@@ -59,6 +59,7 @@ def write_version_file():
 
 requirements = [
     "numpy",
+    "requests",
     pytorch_dep,
 ]
 

setup.py

@@ -29,6 +29,19 @@
 DEFAULT_TEST_DECODER = object()
 
 
+class TestResource(datasets.utils.OnlineResource):
+    def __init__(self, *, dataset_name, dataset_config, **kwargs):
+        super().__init__(**kwargs)
+        self.dataset_name = dataset_name
+        self.dataset_config = dataset_config
+
+    def _download(self, _):
+        raise pytest.UsageError(
+            f"Dataset '{self.dataset_name}' requires the file '{self.file_name}' for {self.dataset_config}, "
+            f"but this file does not exist."
+        )
+
+
 class DatasetMocks:
     def __init__(self):
         self._mock_data_fns = {}
@@ -72,7 +85,7 @@ def _parse_mock_info(self, mock_info, *, name):
             )
         return mock_info
 
-    def _get(self, dataset, config):
+    def _get(self, dataset, config, root):
         name = dataset.info.name
         resources_and_mock_info = self._cache.get((name, config))
         if resources_and_mock_info:
@@ -87,20 +100,12 @@ def _get(self, dataset, config):
                 f"Did you register the mock data function with `@DatasetMocks.register_mock_data_fn`?"
             )
 
-        root = self._tmp_home / name
-        root.mkdir(exist_ok=True)
+        mock_resources = [
+            TestResource(dataset_name=name, dataset_config=config, file_name=resource.file_name)
+            for resource in dataset.resources(config)
+        ]
         mock_info = self._parse_mock_info(fakedata_fn(dataset.info, root, config), name=name)
 
-        mock_resources = []
-        for resource in dataset.resources(config):
-            path = root / resource.file_name
-            if not path.exists() and path.is_file():
-                raise pytest.UsageError(
-                    f"Dataset '{name}' requires the file {path.name} for {config}, but this file does not exist."
-                )
-
-            mock_resources.append(datasets.utils.LocalResource(path))
-
         self._cache[(name, config)] = mock_resources, mock_info
         return mock_resources, mock_info
 
@@ -109,9 +114,13 @@ def load(
     ) -> Tuple[IterDataPipe, Dict[str, Any]]:
         dataset = find(name)
         config = dataset.info.make_config(split=split, **options)
-        resources, mock_info = self._get(dataset, config)
+
+        root = self._tmp_home / name
+        root.mkdir(exist_ok=True)
+        resources, mock_info = self._get(dataset, config, root)
+
         datapipe = dataset._make_datapipe(
-            [resource.to_datapipe() for resource in resources],
+            [resource.load(root) for resource in resources],
             config=config,
             decoder=DEFAULT_DECODER_MAP.get(dataset.info.type) if decoder is DEFAULT_DECODER else decoder,
         )