Add video GPU decoder

setup.py

@@ -427,6 +427,59 @@ def get_extensions():
             )
         )
 
+    # Locating video codec
+    # CUDA_HOME should be set to the cuda root directory.
+    # TORCHVISION_INCLUDE and TORCHVISION_LIBRARY should include the location to
+    # video codec header files and libraries respectively.
+    video_codec_found = (
+        extension is CUDAExtension
+        and CUDA_HOME is not None
+        and any([os.path.exists(os.path.join(folder, "cuviddec.h")) for folder in vision_include])
+        and any([os.path.exists(os.path.join(folder, "nvcuvid.h")) for folder in vision_include])
+        and any([os.path.exists(os.path.join(folder, "libnvcuvid.so")) for folder in library_dirs])
+    )
+
+    print(f"video codec found: {video_codec_found}")
+
+    if (
+        video_codec_found
+        and has_ffmpeg
+        and any([os.path.exists(os.path.join(folder, "libavcodec", "bsf.h")) for folder in ffmpeg_include_dir])
+    ):
+        gpu_decoder_path = os.path.join(extensions_dir, "io", "decoder", "gpu")
+        gpu_decoder_src = glob.glob(os.path.join(gpu_decoder_path, "*.cpp"))
+        cuda_libs = os.path.join(CUDA_HOME, "lib64")
+        cuda_inc = os.path.join(CUDA_HOME, "include")
+
+        ext_modules.append(
+            extension(
+                "torchvision.Decoder",
+                gpu_decoder_src,
+                include_dirs=include_dirs + [gpu_decoder_path] + [cuda_inc] + ffmpeg_include_dir,
+                library_dirs=ffmpeg_library_dir + library_dirs + [cuda_libs],
+                libraries=[
+                    "avcodec",
+                    "avformat",
+                    "avutil",
+                    "swresample",
+                    "swscale",
+                    "nvcuvid",
+                    "cuda",
+                    "cudart",
+                    "z",
+                    "pthread",
+                    "dl",
+                ],
+                extra_compile_args=extra_compile_args,
+            )
+        )
+    else:
+        print(
+            "The installed version of ffmpeg is missing the header file 'bsf.h' which is "
+            "required for GPU video decoding. Please install the latest ffmpeg from conda-forge channel:"
+            " `conda install -c conda-forge ffmpeg`."
+        )
+
     return ext_modules
 
 

test/test_video_gpu_decoder.py

@@ -0,0 +1,41 @@
+import os
+
+import pytest
+import torch
+from torchvision.io import _HAS_VIDEO_DECODER, VideoReader
+
+try:
+    import av
+except ImportError:
+    av = None
+
+VIDEO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "videos")
+
+test_videos = [
+    "RATRACE_wave_f_nm_np1_fr_goo_37.avi",
+    "TrumanShow_wave_f_nm_np1_fr_med_26.avi",
+    "v_SoccerJuggling_g23_c01.avi",
+    "v_SoccerJuggling_g24_c01.avi",
+    "R6llTwEh07w.mp4",
+    "SOX5yA1l24A.mp4",
+    "WUzgd7C1pWA.mp4",
+]
+
+
+@pytest.mark.skipif(_HAS_VIDEO_DECODER is False, reason="Didn't compile with support for gpu decoder")
+class TestVideoGPUDecoder:
+    @pytest.mark.skipif(av is None, reason="PyAV unavailable")
+    def test_frame_reading(self):
+        for test_video in test_videos:
+            full_path = os.path.join(VIDEO_DIR, test_video)
+            decoder = VideoReader(full_path, device="cuda:0")
+            with av.open(full_path) as container:
+                for av_frame in container.decode(container.streams.video[0]):
+                    av_frames = torch.tensor(av_frame.to_ndarray().flatten())
+                    vision_frames = next(decoder)["data"]
+                    mean_delta = torch.mean(torch.abs(av_frames.float() - decoder._reformat(vision_frames).float()))
+                    assert mean_delta < 0.1
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

torchvision/csrc/io/decoder/gpu/README.rst

@@ -0,0 +1,21 @@
+GPU Decoder
+===========
+
+GPU decoder depends on ffmpeg for demuxing, uses NVDECODE APIs from the nvidia-video-codec sdk and uses cuda for processing on gpu. In order to use this, please follow the following steps:
+
+* Download the latest `nvidia-video-codec-sdk <https://developer.nvidia.com/nvidia-video-codec-sdk/download>`_
+* Extract the zipped file.
+* Set TORCHVISION_INCLUDE environment variable to the location of the video codec headers(`nvcuvid.h` and `cuviddec.h`), which would be under `Interface` directory.
+* Set TORCHVISION_LIBRARY environment variable to the location of the video codec library(`libnvcuvid.so`), which would be under `Lib/linux/stubs/x86_64` directory.
+* Install the latest ffmpeg from `conda-forge` channel.
+
+.. code:: bash
+
+    conda install -c conda-forge ffmpeg
+
+* Set CUDA_HOME environment variable to the cuda root directory.
+* Build torchvision from source:
+
+.. code:: bash
+
+    python setup.py install