Merge branch 'main' into patch-1

Author: datumbox

.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

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

packaging/torchvision/meta.yaml

@@ -0,0 +1,57 @@
+# 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. 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
+```

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,
         )

test/builtin_dataset_mocks.py

@@ -211,10 +211,10 @@ def test_default_config(self):
             pytest.param(make_minimal_dataset_info().default_config, None, id="default"),
         ],
     )
-    def test_to_datapipe_config(self, config, kwarg):
+    def test_load_config(self, config, kwarg):
         dataset = self.DatasetMock()
 
-        dataset.to_datapipe("", config=kwarg)
+        dataset.load("", config=kwarg)
 
         dataset.resources.assert_called_with(config)
 
@@ -225,18 +225,19 @@ def test_missing_dependencies(self):
         dependency = "fake_dependency"
         dataset = self.DatasetMock(make_minimal_dataset_info(dependencies=(dependency,)))
         with pytest.raises(ModuleNotFoundError, match=dependency):
-            dataset.to_datapipe("root")
+            dataset.load("root")
 
     def test_resources(self, mocker):
-        resource_mock = mocker.Mock(spec=["to_datapipe"])
+        resource_mock = mocker.Mock(spec=["load"])
         sentinel = object()
-        resource_mock.to_datapipe.return_value = sentinel
+        resource_mock.load.return_value = sentinel
         dataset = self.DatasetMock(resources=[resource_mock])
 
         root = "root"
-        dataset.to_datapipe(root)
+        dataset.load(root)
 
-        resource_mock.to_datapipe.assert_called_with(root)
+        (call_args, _) = resource_mock.load.call_args
+        assert call_args[0] == root
 
         (call_args, _) = dataset._make_datapipe.call_args
         assert call_args[0][0] is sentinel
@@ -245,7 +246,7 @@ def test_decoder(self):
         dataset = self.DatasetMock()
 
         sentinel = object()
-        dataset.to_datapipe("", decoder=sentinel)
+        dataset.load("", decoder=sentinel)
 
         (_, call_kwargs) = dataset._make_datapipe.call_args
         assert call_kwargs["decoder"] is sentinel