Repository: pytorch/vision

Issue:
Implement JPEG decoding via nvjpeg
## 🚀 Feature

We currently provide CPU image decoding functions for PNG and JPEG. It could be interesting to consider using nvjpeg to perform decoding directly on the GPU, as it could benefit certain workloads.

## Motivation

With CUDA 10.0, nvjpeg is [part of the CUDA toolkit](https://developer.nvidia.com/nvjpeg), which means that it would be easily accessible for us to use it.

The API for the decoding function is not yet clear to me, and would require some digging
Ideally, passing a CUDA tensor for `decode_jpeg` should do the decoding on the GPU, although I'm not sure if that's how nvjpeg works internally.



Discussion:
It looks like it requires a host pointer and does some preprocessing on the cpu


> | Parameter |  Input / Output |  Memory  | Description |
> |-------------------------|---|---|----|
> | const unsigned char *data |  Input | Host  | Pointer to the encoded data.
> | nvjpegImage_t *destination |  Input/Output | Host/ Device | Pointer to the structure that describes the output destination. This structure should be on the host (CPU), but the pointers in this structure should be pointing to the device (i.e., GPU) memory. See nvjpegImage_t.

so passing gpu data directly wouldn't work
Yes, that was my understanding as well. We would need to see what potential speed benefits single-image decoding could bring compared to say turbo-jpeg, and if the speed improvements are there what would be a good API for using it. Passing a `device` flag to `decode_jpeg`?

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -315,8 +315,23 @@ def get_extensions():
             image_library += [jpeg_lib]
             image_include += [jpeg_include]
 
+    # Locating nvjpeg
+    # Should be included in CUDA_HOME for CUDA >= 10.1, which is the minimum version we have in the CI
+    nvjpeg_found = (
+        extension is CUDAExtension and
+        CUDA_HOME is not None and
+        os.path.exists(os.path.join(CUDA_HOME, 'include', 'nvjpeg.h'))
+    )
+
+    print('NVJPEG found: {0}'.format(nvjpeg_found))
+    image_macros += [('NVJPEG_FOUND', str(int(nvjpeg_found)))]
+    if nvjpeg_found:
+        print('Building torchvision with NVJPEG image support')
+        image_link_flags.append('nvjpeg')
+
     image_path = os.path.join(extensions_dir, 'io', 'image')
-    image_src = glob.glob(os.path.join(image_path, '*.cpp')) + glob.glob(os.path.join(image_path, 'cpu', '*.cpp'))
+    image_src = (glob.glob(os.path.join(image_path, '*.cpp')) + glob.glob(os.path.join(image_path, 'cpu', '*.cpp'))
+                 + glob.glob(os.path.join(image_path, 'cuda', '*.cpp')))
 
     if png_found or jpeg_found:
         ext_modules.append(extension(
diff --git a/torchvision/io/image.py b/torchvision/io/image.py
--- a/torchvision/io/image.py
+++ b/torchvision/io/image.py
@@ -148,7 +148,8 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6):
     write_file(filename, output)
 
 
-def decode_jpeg(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
+def decode_jpeg(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED,
+                device: str = 'cpu') -> torch.Tensor:
     """
     Decodes a JPEG image into a 3 dimensional RGB Tensor.
     Optionally converts the image to the desired format.
@@ -156,16 +157,25 @@ def decode_jpeg(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANG
 
     Args:
         input (Tensor[1]): a one dimensional uint8 tensor containing
-            the raw bytes of the JPEG image.
+            the raw bytes of the JPEG image. This tensor must be on CPU,
+            regardless of the ``device`` parameter.
         mode (ImageReadMode): the read mode used for optionally
             converting the image. Default: `ImageReadMode.UNCHANGED`.
             See `ImageReadMode` class for more information on various
             available modes.
+        device (str or torch.device): The device on which the decoded image will
+            be stored. If a cuda device is specified, the image will be decoded
+            with `nvjpeg <https://developer.nvidia.com/nvjpeg>`_. This is only
+            supported for CUDA version >= 10.1
 
     Returns:
         output (Tensor[image_channels, image_height, image_width])
     """
-    output = torch.ops.image.decode_jpeg(input, mode.value)
+    device = torch.device(device)
+    if device.type == 'cuda':
+        output = torch.ops.image.decode_jpeg_cuda(input, mode.value, device)
+    else:
+        output = torch.ops.image.decode_jpeg(input, mode.value)
     return output
 
 


Test patch:
diff --git a/test/common_utils.py b/test/common_utils.py
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -24,6 +24,9 @@
 PY39_SEGFAULT_SKIP_MSG = "Segmentation fault with Python 3.9, see https://github.com/pytorch/vision/issues/3367"
 PY39_SKIP = unittest.skipIf(IS_PY39, PY39_SEGFAULT_SKIP_MSG)
 IN_CIRCLE_CI = os.getenv("CIRCLECI", False) == 'true'
+IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None
+IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1"
+CUDA_NOT_AVAILABLE_MSG = 'CUDA device not available'
 
 
 @contextlib.contextmanager
@@ -407,11 +410,8 @@ def call_args_to_kwargs_only(call_args, *callable_or_arg_names):
 
 def cpu_and_gpu():
     import pytest  # noqa
-    # ignore CPU tests in RE as they're already covered by another contbuild
-    IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None
-    IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1"
-    CUDA_NOT_AVAILABLE_MSG = 'CUDA device not available'
 
+    # ignore CPU tests in RE as they're already covered by another contbuild
     devices = [] if IN_RE_WORKER else ['cpu']
 
     if torch.cuda.is_available():
@@ -427,3 +427,17 @@ def cpu_and_gpu():
     devices.append(pytest.param('cuda', marks=cuda_marks))
 
     return devices
+
+
+def needs_cuda(test_func):
+    import pytest  # noqa
+
+    if IN_FBCODE and not IN_RE_WORKER:
+        # We don't want to skip in fbcode, so we just don't collect
+        # TODO: slightly more robust way would be to detect if we're in a sandcastle instance
+        # so that the test will still be collected (and skipped) in the devvms.
+        return pytest.mark.dont_collect(test_func)
+    elif torch.cuda.is_available():
+        return test_func
+    else:
+        return pytest.mark.skip(reason=CUDA_NOT_AVAILABLE_MSG)(test_func)
diff --git a/test/test_image.py b/test/test_image.py
--- a/test/test_image.py
+++ b/test/test_image.py
@@ -3,10 +3,11 @@
 import os
 import unittest
 
+import pytest
 import numpy as np
 import torch
 from PIL import Image
-from common_utils import get_tmp_dir
+from common_utils import get_tmp_dir, needs_cuda
 
 from torchvision.io.image import (
     decode_png, decode_jpeg, encode_jpeg, write_jpeg, decode_image, read_file,
@@ -278,5 +279,43 @@ def test_write_file_non_ascii(self):
             os.unlink(fpath)
 
 
+@needs_cuda
+@pytest.mark.parametrize('mode', [ImageReadMode.UNCHANGED, ImageReadMode.GRAY, ImageReadMode.RGB])
+@pytest.mark.parametrize('img_path', get_images(IMAGE_ROOT, ".jpg"))
+@pytest.mark.parametrize('scripted', (False, True))
+def test_decode_jpeg_cuda(mode, img_path, scripted):
+    if 'cmyk' in img_path:
+        pytest.xfail("Decoding a CMYK jpeg isn't supported")
+    tester = ImageTester()
+    data = read_file(img_path)
+    img = decode_image(data, mode=mode)
+    f = torch.jit.script(decode_jpeg) if scripted else decode_jpeg
+    img_nvjpeg = f(data, mode=mode, device='cuda')
+
+    # Some difference expected between jpeg implementations
+    tester.assertTrue((img.float() - img_nvjpeg.cpu().float()).abs().mean() < 2)
+
+
+@needs_cuda
+@pytest.mark.parametrize('cuda_device', ('cuda', 'cuda:0', torch.device('cuda')))
+def test_decode_jpeg_cuda_device_param(cuda_device):
+    """Make sure we can pass a string or a torch.device as device param"""
+    data = read_file(next(get_images(IMAGE_ROOT, ".jpg")))
+    decode_jpeg(data, device=cuda_device)
+
+
+@needs_cuda
+def test_decode_jpeg_cuda_errors():
+    data = read_file(next(get_images(IMAGE_ROOT, ".jpg")))
+    with pytest.raises(RuntimeError, match="Expected a non empty 1-dimensional tensor"):
+        decode_jpeg(data.reshape(-1, 1), device='cuda')
+    with pytest.raises(RuntimeError, match="input tensor must be on CPU"):
+        decode_jpeg(data.to('cuda'), device='cuda')
+    with pytest.raises(RuntimeError, match="Expected a torch.uint8 tensor"):
+        decode_jpeg(data.to(torch.float), device='cuda')
+    with pytest.raises(RuntimeError, match="Expected a cuda device"):
+        torch.ops.image.decode_jpeg_cuda(data, ImageReadMode.UNCHANGED.value, 'cpu')
+
+
 if __name__ == '__main__':
     unittest.main()