Merge branch 'main' into datasets/video-decoding

.circleci/config.yml.in

@@ -848,7 +848,7 @@ jobs:
     executor:
       name: windows-gpu
     environment:
-      CUDA_VERSION: "11.1"
+      CUDA_VERSION: "11.3"
       PYTHON_VERSION: << parameters.python_version >>
     steps:
       - checkout

.circleci/regenerate.py

@@ -21,7 +21,7 @@
 from jinja2 import select_autoescape
 
 
-PYTHON_VERSIONS = ["3.7", "3.8", "3.9"]
+PYTHON_VERSIONS = ["3.7", "3.8", "3.9", "3.10"]
 
 RC_PATTERN = r"/v[0-9]+(\.[0-9]+)*-rc[0-9]+/"
 
@@ -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", "cu115", "rocm4.3.1", "rocm4.5.2"],
-                "win": ["cpu", "cu111", "cu113", "cu115"],
+                "linux": ["cpu", "cu102", "cu113", "cu115", "rocm4.3.1", "rocm4.5.2"],
+                "win": ["cpu", "cu113", "cu115"],
                 "macos": ["cpu"],
             }
             cu_versions = cu_versions_dict[os_type]

.circleci/unittest/windows/scripts/environment.yml

@@ -9,10 +9,11 @@ dependencies:
   - libpng
   - jpeg
   - ca-certificates
-  - h5py
+  - hdf5
   - pip:
     - future
     - pillow >=5.3.0, !=8.3.*
     - scipy
     - av
     - dataclasses
+    - h5py

.circleci/unittest/windows/scripts/install.sh

@@ -28,7 +28,7 @@ else
 fi
 
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"
-conda install -y -c "pytorch-${UPLOAD_CHANNEL}" -c conda-forge "pytorch-${UPLOAD_CHANNEL}"::pytorch[build="*${version}*"] "${cudatoolkit}" pytest
+conda install -y -c "pytorch-${UPLOAD_CHANNEL}" -c nvidia "pytorch-${UPLOAD_CHANNEL}"::pytorch[build="*${version}*"] "${cudatoolkit}" pytest
 
 torch_cuda=$(python -c "import torch; print(torch.cuda.is_available())")
 echo torch.cuda.is_available is $torch_cuda

.pre-commit-config.yaml

@@ -10,15 +10,8 @@ repos:
         args: [--fix=lf]
       - id: end-of-file-fixer
 
-  # - repo: https://github.com/asottile/pyupgrade
-  #   rev: v2.29.0
-  #   hooks:
-  #     - id: pyupgrade
-  #       args: [--py37-plus]
-  #       name: Upgrade code
-
   - repo: https://github.com/omnilib/ufmt
-    rev: v1.3.0
+    rev: v1.3.2
     hooks:
       - id: ufmt
         additional_dependencies:

docs/source/ops.rst

@@ -20,6 +20,8 @@ Operators
     box_iou
     clip_boxes_to_image
     deform_conv2d
+    drop_block2d
+    drop_block3d
     generalized_box_iou
     generalized_box_iou_loss
     masks_to_boxes
@@ -45,5 +47,8 @@ Operators
     FeaturePyramidNetwork
     StochasticDepth
     FrozenBatchNorm2d
-    ConvNormActivation
+    Conv2dNormActivation
+    Conv3dNormActivation
     SqueezeExcitation
+    DropBlock2d
+    DropBlock3d

docs/source/transforms.rst

@@ -198,6 +198,7 @@ The new transform can be used standalone or mixed-and-matched with existing tran
     AutoAugment
     RandAugment
     TrivialAugmentWide
+    AugMix
 
 .. _functional_transforms:
 

gallery/plot_transforms.py

@@ -263,6 +263,14 @@ def plot(imgs, with_orig=True, row_title=None, **imshow_kwargs):
 imgs = [augmenter(orig_img) for _ in range(4)]
 plot(imgs)
 
+####################################
+# AugMix
+# ~~~~~~
+# The :class:`~torchvision.transforms.AugMix` transform automatically augments the data.
+augmenter = T.AugMix()
+imgs = [augmenter(orig_img) for _ in range(4)]
+plot(imgs)
+
 ####################################
 # Randomly-applied transforms
 # ---------------------------

packaging/pkg_helpers.bash

@@ -205,6 +205,7 @@ setup_wheel_python() {
       3.7) python_abi=cp37-cp37m ;;
       3.8) python_abi=cp38-cp38 ;;
       3.9) python_abi=cp39-cp39 ;;
+      3.10) python_abi=cp310-cp310 ;;
       *)
         echo "Unrecognized PYTHON_VERSION=$PYTHON_VERSION"
         exit 1

references/classification/presets.py

@@ -22,6 +22,8 @@ def __init__(
                 trans.append(autoaugment.RandAugment(interpolation=interpolation))
             elif auto_augment_policy == "ta_wide":
                 trans.append(autoaugment.TrivialAugmentWide(interpolation=interpolation))
+            elif auto_augment_policy == "augmix":
+                trans.append(autoaugment.AugMix(interpolation=interpolation))
             else:
                 aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
                 trans.append(autoaugment.AutoAugment(policy=aa_policy, interpolation=interpolation))

references/classification/utils.py

@@ -166,17 +166,7 @@ def __init__(self, model, decay, device="cpu"):
         def ema_avg(avg_model_param, model_param, num_averaged):
             return decay * avg_model_param + (1 - decay) * model_param
 
-        super().__init__(model, device, ema_avg)
-
-    def update_parameters(self, model):
-        for p_swa, p_model in zip(self.module.state_dict().values(), model.state_dict().values()):
-            device = p_swa.device
-            p_model_ = p_model.detach().to(device)
-            if self.n_averaged == 0:
-                p_swa.detach().copy_(p_model_)
-            else:
-                p_swa.detach().copy_(self.avg_fn(p_swa.detach(), p_model_, self.n_averaged.to(device)))
-        self.n_averaged += 1
+        super().__init__(model, device, ema_avg, use_buffers=True)
 
 
 def accuracy(output, target, topk=(1,)):
@@ -274,6 +264,7 @@ def init_distributed_mode(args):
     torch.distributed.init_process_group(
         backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
     )
+    torch.distributed.barrier()
     setup_for_distributed(args.rank == 0)
 
 

references/detection/transforms.py

@@ -4,7 +4,7 @@
 import torchvision
 from torch import nn, Tensor
 from torchvision.transforms import functional as F
-from torchvision.transforms import transforms as T
+from torchvision.transforms import transforms as T, InterpolationMode
 
 
 def _flip_coco_person_keypoints(kps, width):
@@ -282,3 +282,55 @@ def forward(
                 image = F.to_pil_image(image)
 
         return image, target
+
+
+class ScaleJitter(nn.Module):
+    """Randomly resizes the image and its bounding boxes  within the specified scale range.
+    The class implements the Scale Jitter augmentation as described in the paper
+    `"Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation" <https://arxiv.org/abs/2012.07177>`_.
+
+    Args:
+        target_size (tuple of ints): The target size for the transform provided in (height, weight) format.
+        scale_range (tuple of ints): scaling factor interval, e.g (a, b), then scale is randomly sampled from the
+            range a <= scale <= b.
+        interpolation (InterpolationMode): Desired interpolation enum defined by
+            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
+    """
+
+    def __init__(
+        self,
+        target_size: Tuple[int, int],
+        scale_range: Tuple[float, float] = (0.1, 2.0),
+        interpolation: InterpolationMode = InterpolationMode.BILINEAR,
+    ):
+        super().__init__()
+        self.target_size = target_size
+        self.scale_range = scale_range
+        self.interpolation = interpolation
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        orig_width, orig_height = F.get_image_size(image)
+
+        r = self.scale_range[0] + torch.rand(1) * (self.scale_range[1] - self.scale_range[0])
+        new_width = int(self.target_size[1] * r)
+        new_height = int(self.target_size[0] * r)
+
+        image = F.resize(image, [new_height, new_width], interpolation=self.interpolation)
+
+        if target is not None:
+            target["boxes"][:, 0::2] *= new_width / orig_width
+            target["boxes"][:, 1::2] *= new_height / orig_height
+            if "masks" in target:
+                target["masks"] = F.resize(
+                    target["masks"], [new_height, new_width], interpolation=InterpolationMode.NEAREST
+                )
+
+        return image, target

references/optical_flow/utils.py

@@ -267,6 +267,7 @@ def setup_ddp(args):
         world_size=args.world_size,
         init_method=args.dist_url,
     )
+    torch.distributed.barrier()
 
 
 def reduce_across_processes(val):

references/segmentation/utils.py

@@ -291,4 +291,5 @@ def init_distributed_mode(args):
     torch.distributed.init_process_group(
         backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
     )
+    torch.distributed.barrier()
     setup_for_distributed(args.rank == 0)

references/video_classification/utils.py

@@ -250,4 +250,5 @@ def init_distributed_mode(args):
     torch.distributed.init_process_group(
         backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
     )
+    torch.distributed.barrier()
     setup_for_distributed(args.rank == 0)

scripts/release_notes/classify_prs.py

@@ -1,13 +1,18 @@
 # In[1]:
 
-
+# imports and set configuration
 import pandas as pd
+from retrieve_prs_data import run
 
+exclude_prototype = True
+data_filename = "10.0_to_11.0-rc2.json"
+previous_release = "v10.0"
+current_release = "v11.0-rc2"
 
 # In[2]:
 
 
-df = pd.read_json("10.0_to_11.0-rc2.json").T
+df = pd.read_json(data_filename).T
 df.tail()
 
 
@@ -76,6 +81,8 @@
 def format_prs(mod_df):
     out = []
     for idx, row in mod_df.iterrows():
+        if exclude_prototype and row["prototype"]:
+            continue
         modules = idx
         # Put "documentation" and "tests" first for sorting to be dece
         for last_module in ("documentation", "tests"):
@@ -119,3 +126,13 @@ def format_prs(mod_df):
 # Missing PRs are these ones... classify them manually
 missing_prs = pd.concat([mod_df, included_prs]).drop_duplicates(subset="pr_number", keep=False)
 print(format_prs(missing_prs))
+
+# In[12]:
+
+# Generate list of contributors
+print()
+print("## Contributors")
+
+command_to_run = f"{{ git shortlog -s {previous_release}..{current_release} | cut -f2- & git log -s {previous_release}..{current_release} | grep Co-authored | cut -f2- -d: | cut -f1 -d\\< | sed 's/^ *//;s/ *$//' ; }} | sort --ignore-case | uniq | tr '\\n' ';' | sed 's/;/, /g;s/, $//' | fold -s"
+rc, output, err = run(command_to_run)
+print(output)

test/test_ops.py

@@ -2,6 +2,7 @@
 import os
 from abc import ABC, abstractmethod
 from functools import lru_cache
+from itertools import product
 from typing import Callable, List, Tuple
 
 import numpy as np
@@ -57,6 +58,16 @@ def forward(self, a):
         self.layer(a)
 
 
+class DropBlockWrapper(nn.Module):
+    def __init__(self, obj):
+        super().__init__()
+        self.layer = obj
+        self.n_inputs = 1
+
+    def forward(self, a):
+        self.layer(a)
+
+
 class RoIOpTester(ABC):
     dtype = torch.float64
 
@@ -1357,5 +1368,87 @@ def test_split_normalization_params(self, norm_layer):
         assert len(params[1]) == 82
 
 
+class TestDropBlock:
+    @pytest.mark.parametrize("seed", range(10))
+    @pytest.mark.parametrize("dim", [2, 3])
+    @pytest.mark.parametrize("p", [0, 0.5])
+    @pytest.mark.parametrize("block_size", [5, 11])
+    @pytest.mark.parametrize("inplace", [True, False])
+    def test_drop_block(self, seed, dim, p, block_size, inplace):
+        torch.manual_seed(seed)
+        batch_size = 5
+        channels = 3
+        height = 11
+        width = height
+        depth = height
+        if dim == 2:
+            x = torch.ones(size=(batch_size, channels, height, width))
+            layer = ops.DropBlock2d(p=p, block_size=block_size, inplace=inplace)
+            feature_size = height * width
+        elif dim == 3:
+            x = torch.ones(size=(batch_size, channels, depth, height, width))
+            layer = ops.DropBlock3d(p=p, block_size=block_size, inplace=inplace)
+            feature_size = depth * height * width
+        layer.__repr__()
+
+        out = layer(x)
+        if p == 0:
+            assert out.equal(x)
+        if block_size == height:
+            for b, c in product(range(batch_size), range(channels)):
+                assert out[b, c].count_nonzero() in (0, feature_size)
+
+    @pytest.mark.parametrize("seed", range(10))
+    @pytest.mark.parametrize("dim", [2, 3])
+    @pytest.mark.parametrize("p", [0.1, 0.2])
+    @pytest.mark.parametrize("block_size", [3])
+    @pytest.mark.parametrize("inplace", [False])
+    def test_drop_block_random(self, seed, dim, p, block_size, inplace):
+        torch.manual_seed(seed)
+        batch_size = 5
+        channels = 3
+        height = 11
+        width = height
+        depth = height
+        if dim == 2:
+            x = torch.ones(size=(batch_size, channels, height, width))
+            layer = ops.DropBlock2d(p=p, block_size=block_size, inplace=inplace)
+        elif dim == 3:
+            x = torch.ones(size=(batch_size, channels, depth, height, width))
+            layer = ops.DropBlock3d(p=p, block_size=block_size, inplace=inplace)
+
+        trials = 250
+        num_samples = 0
+        counts = 0
+        cell_numel = torch.tensor(x.shape).prod()
+        for _ in range(trials):
+            with torch.no_grad():
+                out = layer(x)
+            non_zero_count = out.nonzero().size(0)
+            counts += cell_numel - non_zero_count
+            num_samples += cell_numel
+
+        assert abs(p - counts / num_samples) / p < 0.15
+
+    def make_obj(self, dim, p, block_size, inplace, wrap=False):
+        if dim == 2:
+            obj = ops.DropBlock2d(p, block_size, inplace)
+        elif dim == 3:
+            obj = ops.DropBlock3d(p, block_size, inplace)
+        return DropBlockWrapper(obj) if wrap else obj
+
+    @pytest.mark.parametrize("dim", (2, 3))
+    @pytest.mark.parametrize("p", [0, 1])
+    @pytest.mark.parametrize("block_size", [5, 7])
+    @pytest.mark.parametrize("inplace", [True, False])
+    def test_is_leaf_node(self, dim, p, block_size, inplace):
+        op_obj = self.make_obj(dim, p, block_size, inplace, wrap=True)
+        graph_node_names = get_graph_node_names(op_obj)
+
+        assert len(graph_node_names) == 2
+        assert len(graph_node_names[0]) == len(graph_node_names[1])
+        assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
+
+
 if __name__ == "__main__":
     pytest.main([__file__])