Repository: pytorch/vision
Pull request: Adding Uniform temporal Subsampling for Video
Resolves #6768

Adding temporal sampling Transforms similar to the ones on [PyTorch Video](https://github.com/facebookresearch/pytorchvideo/blob/a0a131e/pytorchvideo/transforms/functional.py#L19):
- [x] Dispatcher + Kernel
- [x] Transform Class
- [x] Tests

cc @vfdev-5 @bjuncek @pmeier

Problem statement:
[prototype] Add temporal Jittering and Permute transforms and kernels
### 🚀 The feature

Though the sampling of the temporal dimension can be handled efficiently on the VideoClipper (work done by @YosuaMichael), some users still want to perform jittering on the temporal dimension as part of the transforms. This can be achieved by uniformly sampling the indeces of the available samples and selecting the appropriate frames.

Similarly we need an easy transform for permuting channel order.

### Motivation, pitch

Add a Temporal Sampler and a Permute method as transform along with the necessary kernels.

### Alternatives

_No response_

### Additional context

_No response_

cc @vfdev-5 @bjuncek @pmeier

Hints from issue discussion:

Commit message: add more tests for uniform_temporal_subsample

Patch to review:
diff --git a/.github/workflows/prototype-tests.yml b/.github/workflows/prototype-tests.yml
index 5e9ca360d08..21b112e4a79 100644
--- a/.github/workflows/prototype-tests.yml
+++ b/.github/workflows/prototype-tests.yml
@@ -34,7 +34,7 @@ jobs:
         run: pip install --progress-bar=off --no-build-isolation --editable .
 
       - name: Install other prototype dependencies
-        run: pip install --progress-bar=off scipy pycocotools h5py iopath
+        run: pip install --progress-bar=off scipy pycocotools h5py iopath pytorchvideo
 
       - name: Install test requirements
         run: pip install --progress-bar=off pytest pytest-mock pytest-cov
diff --git a/test/prototype_transforms_dispatcher_infos.py b/test/prototype_transforms_dispatcher_infos.py
index fab4b3583d5..6aa7eb248ca 100644
--- a/test/prototype_transforms_dispatcher_infos.py
+++ b/test/prototype_transforms_dispatcher_infos.py
@@ -426,4 +426,13 @@ def xfail_all_tests(*, reason, condition):
             skip_dispatch_feature,
         ],
     ),
+    DispatcherInfo(
+        F.uniform_temporal_subsample,
+        kernels={
+            features.Video: F.uniform_temporal_subsample_video,
+        },
+        test_marks=[
+            skip_dispatch_feature,
+        ],
+    ),
 ]
diff --git a/test/prototype_transforms_kernel_infos.py b/test/prototype_transforms_kernel_infos.py
index 77a77444ba7..d5ffdd86b38 100644
--- a/test/prototype_transforms_kernel_infos.py
+++ b/test/prototype_transforms_kernel_infos.py
@@ -1,5 +1,6 @@
 import decimal
 import functools
+import importlib
 import itertools
 import math
 import re
@@ -132,6 +133,23 @@ def xfail_all_tests(*, reason, condition):
     ]
 
 
+def package_is_available(name):
+    try:
+        importlib.import_module(name)
+        # FIXME
+        return False
+    except ModuleNotFoundError:
+        return False
+
+
+def skip_on_missing_third_party_reference(name):
+    return TestMark(
+        ("TestKernels", "test_against_reference"),
+        pytest.mark.skip(f"Third party package '{name}' is not available"),
+        condition=lambda args_kwargs: package_is_available(name),
+    )
+
+
 KERNEL_INFOS = []
 
 
@@ -2100,3 +2118,39 @@ def sample_inputs_convert_dtype_video():
         ),
     ]
 )
+
+
+def sample_inputs_uniform_temporal_subsample_video():
+    for video_loader in make_video_loaders(sizes=["random"], num_frames=[4]):
+        for temporal_dim in [-4, len(video_loader.shape) - 4]:
+            yield ArgsKwargs(video_loader, num_samples=2, temporal_dim=temporal_dim)
+
+
+def reference_uniform_temporal_subsample_video(video, num_samples, temporal_dim=-4):
+    from pytorchvideo.transforms.functional import uniform_temporal_subsample
+
+    return uniform_temporal_subsample(video, num_samples=num_samples, temporal_dim=temporal_dim)
+
+
+def reference_inputs_uniform_temporal_subsample_video():
+    for video_loader in make_video_loaders(sizes=["random"], color_spaces=[features.ColorSpace.RGB], num_frames=[10]):
+        for num_frames in range(1, 11):
+            yield ArgsKwargs(video_loader, num_frames)
+
+
+KERNEL_INFOS.append(
+    KernelInfo(
+        F.uniform_temporal_subsample_video,
+        sample_inputs_fn=sample_inputs_uniform_temporal_subsample_video,
+        reference_fn=reference_uniform_temporal_subsample_video,
+        reference_inputs_fn=reference_inputs_uniform_temporal_subsample_video,
+        test_marks=[
+            skip_on_missing_third_party_reference("pytorchvideo"),
+            TestMark(
+                ("TestKernels", "test_batched_vs_single"),
+                pytest.mark.skip("Positive `temporal_dim` arguments are not equivalent for batched and single inputs"),
+                condition=lambda args_kwargs: args_kwargs.kwargs.get("temporal_dim") >= 0,
+            ),
+        ],
+    )
+)
diff --git a/test/test_prototype_transforms.py b/test/test_prototype_transforms.py
index 1a4a098dbab..18b362bb6f0 100644
--- a/test/test_prototype_transforms.py
+++ b/test/test_prototype_transforms.py
@@ -1899,3 +1899,22 @@ def test_transpose_dimensions(dims):
             assert type(transformed_value) == torch.Tensor
         else:
             assert transformed_value is value
+
+
+class TestUniformTemporalSubsample:
+    @pytest.mark.parametrize(
+        "inpt",
+        [
+            torch.zeros(10, 3, 8, 8),
+            torch.zeros(1, 10, 3, 8, 8),
+            features.Video(torch.zeros(1, 10, 3, 8, 8)),
+        ],
+    )
+    def test__transform(self, inpt):
+        num_samples = 5
+        transform = transforms.UniformTemporalSubsample(num_samples)
+
+        output = transform(inpt)
+        assert isinstance(output, inpt.__class__)
+        assert output.shape[-4] == num_samples
+        assert output.dtype == inpt.dtype
diff --git a/test/test_prototype_transforms_functional.py b/test/test_prototype_transforms_functional.py
index cad5c204af8..f0b8511e18d 100644
--- a/test/test_prototype_transforms_functional.py
+++ b/test/test_prototype_transforms_functional.py
@@ -242,6 +242,7 @@ def test_scripted_smoke(self, info, args_kwargs, device):
             F.get_num_frames,
             F.get_spatial_size,
             F.rgb_to_grayscale,
+            F.uniform_temporal_subsample,
         ],
         ids=lambda dispatcher: dispatcher.__name__,
     )
@@ -1060,3 +1061,13 @@ def test_equalize_image_tensor_edge_cases():
     inpt[..., 100:, 100:] = 1
     output = F.equalize_image_tensor(inpt)
     assert output.unique().tolist() == [0, 255]
+
+
+@pytest.mark.parametrize("device", cpu_and_gpu())
+def test_correctness_uniform_temporal_subsample(device):
+    video = torch.arange(10, device=device)[:, None, None, None].expand(-1, 3, 8, 8)
+    out_video = F.uniform_temporal_subsample(video, 5)
+    assert out_video.unique().tolist() == [0, 2, 4, 6, 9]
+
+    out_video = F.uniform_temporal_subsample(video, 8)
+    assert out_video.unique().tolist() == [0, 1, 2, 3, 5, 6, 7, 9]
\ No newline at end of file
diff --git a/torchvision/prototype/transforms/__init__.py b/torchvision/prototype/transforms/__init__.py
index 099c30c9c6c..2f9bd76d4ed 100644
--- a/torchvision/prototype/transforms/__init__.py
+++ b/torchvision/prototype/transforms/__init__.py
@@ -51,6 +51,7 @@
     ToDtype,
     TransposeDimensions,
 )
+from ._temporal import UniformTemporalSubsample
 from ._type_conversion import DecodeImage, LabelToOneHot, PILToTensor, ToImagePIL, ToImageTensor, ToPILImage
 
 from ._deprecated import Grayscale, RandomGrayscale, ToTensor  # usort: skip
diff --git a/torchvision/prototype/transforms/_temporal.py b/torchvision/prototype/transforms/_temporal.py
new file mode 100644
index 00000000000..46293c25131
--- /dev/null
+++ b/torchvision/prototype/transforms/_temporal.py
@@ -0,0 +1,16 @@
+from typing import Any, Dict
+
+from torchvision.prototype import features
+from torchvision.prototype.transforms import functional as F, Transform
+
+
+class UniformTemporalSubsample(Transform):
+    _transformed_types = (features.is_simple_tensor, features.Video)
+
+    def __init__(self, num_samples: int, temporal_dim: int = -4):
+        super().__init__()
+        self.num_samples = num_samples
+        self.temporal_dim = temporal_dim
+
+    def _transform(self, inpt: features.VideoType, params: Dict[str, Any]) -> features.VideoType:
+        return F.uniform_temporal_subsample(inpt, self.num_samples, temporal_dim=self.temporal_dim)
diff --git a/torchvision/prototype/transforms/functional/__init__.py b/torchvision/prototype/transforms/functional/__init__.py
index 7e520c98691..0b49c53b5fb 100644
--- a/torchvision/prototype/transforms/functional/__init__.py
+++ b/torchvision/prototype/transforms/functional/__init__.py
@@ -165,6 +165,7 @@
     normalize_image_tensor,
     normalize_video,
 )
+from ._temporal import uniform_temporal_subsample, uniform_temporal_subsample_video
 from ._type_conversion import (
     decode_image_with_pil,
     decode_video_with_av,
diff --git a/torchvision/prototype/transforms/functional/_temporal.py b/torchvision/prototype/transforms/functional/_temporal.py
new file mode 100644
index 00000000000..e2cf0bb1f06
--- /dev/null
+++ b/torchvision/prototype/transforms/functional/_temporal.py
@@ -0,0 +1,24 @@
+import torch
+
+from torchvision.prototype import features
+
+
+def uniform_temporal_subsample_video(video: torch.Tensor, num_samples: int, temporal_dim: int = -4) -> torch.Tensor:
+    # Reference: https://github.com/facebookresearch/pytorchvideo/blob/a0a131e/pytorchvideo/transforms/functional.py#L19
+    t_max = video.shape[temporal_dim] - 1
+    indices = torch.linspace(0, t_max, num_samples, device=video.device).long()
+    return torch.index_select(video, temporal_dim, indices)
+
+
+def uniform_temporal_subsample(
+    inpt: features.VideoTypeJIT, num_samples: int, temporal_dim: int = -4
+) -> features.VideoTypeJIT:
+    if isinstance(inpt, torch.Tensor) and (torch.jit.is_scripting() or not isinstance(inpt, features.Video)):
+        return uniform_temporal_subsample_video(inpt, num_samples, temporal_dim=temporal_dim)
+    else:  # isinstance(inpt, features.Video)
+        if temporal_dim != -4 or inpt.ndim - 4 != temporal_dim:
+            raise ValueError("Video inputs must have temporal_dim equivalent to -4")
+        output = uniform_temporal_subsample_video(
+            inpt.as_subclass(torch.Tensor), num_samples, temporal_dim=temporal_dim
+        )
+        return features.Video.wrap_like(inpt, output)


Code review comments:
diff_hunk: @@ -2100,3 +2118,39 @@ def sample_inputs_convert_dtype_video():
         ),
     ]
 )
+
+
+def sample_inputs_uniform_temporal_subsample_video():
+    for video_loader in make_video_loaders(sizes=["random"], num_frames=[4]):
+        for temporal_dim in [-4, len(video_loader.shape) - 4]:
+            yield ArgsKwargs(video_loader, num_samples=2, temporal_dim=temporal_dim)
+
+
+def reference_uniform_temporal_subsample_video(video, num_samples, temporal_dim=-4):
+    from pytorchvideo.transforms.functional import uniform_temporal_subsample
+
+    return uniform_temporal_subsample(video, num_samples=num_samples, temporal_dim=temporal_dim)
path: test/prototype_transforms_kernel_infos.py
text: @user1:
If there is a good reference function, we can optionally define a `reference_fn`

https://github.com/pytorch/vision/blob/9c112935abe400222cca8f9fbc2d8386e0f25e80/test/prototype_transforms_kernel_infos.py#L45-L49

In our case we have the perfect reference as we want to mirror the function of `pytorchvideo`. We import it lazily here since we shouldn't expect it to be available when this module is collected. I'll explain what else we need to do to skip this test in case we don't have `pytorchvideo` below.

diff_hunk: @@ -34,7 +34,7 @@ jobs:
         run: pip install --progress-bar=off --no-build-isolation --editable .
 
       - name: Install other prototype dependencies
-        run: pip install --progress-bar=off scipy pycocotools h5py iopath
+        run: pip install --progress-bar=off scipy pycocotools h5py iopath pytorchvideo
path: .github/workflows/prototype-tests.yml
text: @user1:
Since we want to run the reference test in CI, we add `pytorchvideo` to the third party packages we install. Since it is a pure Python library, there is no need to install it with the other PyTorch packages like `torch` or `torchdata` above.

@author:
Let's remove the dependency. Adding this is a bit of an overkill; I understand that this is on the tests only but pytorchvideo is a research library which breaks from time to time and also has a cyclical dependency on TorchVision. If it's absolutely needed to have a reference implementation, we could copy it over on our references scripts.

@user1:
It is optional as explained above, but I don't see a reason to not have it. I mean things like https://github.com/pytorch/vision/pull/6812#discussion_r1002437589 make it clear that there are subtle things that we need to watch out for when refactoring the kernel. Why not encode the "right" behavior in an unoptimized reference function?

@user1:
I've removed `pytorchvideo` in 82c0aeebd636c8bb75977e007b9bfb340246f5b8 and copy-pasted the implementation into the reference function.