dusty-nv / jetson-containers

Machine Learning Containers for NVIDIA Jetson and JetPack-L4T
MIT License
1.9k stars 416 forks source link

RuntimeError: decode_jpeg_cuda: torchvision not compiled with nvJPEG support on Jetson AGX Orin Development Kit (Jetpack 5.1.2) #446

Closed udit7395 closed 3 months ago

udit7395 commented 3 months ago

torchvision supports decoding jpeg image via nvjpeg but I get an error when using it on Jetson AGX Orin development kit

Code

import torch
import torchvision
device = torch.device("cuda")
img_path = "path_to_your_image"
img_cpu = torchvision.io.read_file(img_path)
img_gpu = torchvision.io.decode_jpeg(img_cpu, device=device)

Error

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.8/dist-packages/torchvision/io/image.py", line 168, in decode_jpeg
    output = torch.ops.image.decode_jpeg_cuda(input, mode.value, device)
  File "/usr/local/lib/python3.8/dist-packages/torch/_ops.py", line 502, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: decode_jpeg_cuda: torchvision not compiled with nvJPEG support

Container Used: dustynv/l4t-pytorch:r35.4.1

I suspect while building torchvision, it was unable to find nvjpeg. This might be because it is trying to look for nvjpeg in CUDA_HOME (/usr/local/cuda) which is true of x86 based systems but not for Jetson's I guess. I ran the find command to find "nvjpeg.h" but was unable to find its location on the Orin. Could you help with this issue? I would like build torchvision with nvJPEG as it helps loads image faster instead of loading image via CPU using opencv and then moving to GPU

Snippet of code used to build torchvision to find nvjpeg

    nvjpeg_found = (
        extension is CUDAExtension
        and CUDA_HOME is not None
        and os.path.exists(os.path.join(CUDA_HOME, "include", "nvjpeg.h"))
    )
dusty-nv commented 3 months ago

@udit7395 I think the issue is that the Jetson's hardware JPEG decoder uses a different API (either through GStreamer or the Jetson Multimedia API), and doesn't use that same nvjpeg.h interface. I can't find nvjpeg.h on my Jetson's either...

udit7395 commented 3 months ago

That makes sense, thanks @dusty-nv. I don't think there is anything that can be done at the container level, so please feel free to close this issue.

For future reference, I raised an issue on the Jetson forums about this.