Torchvision decode_jpeg memory leak

[igrekun]

🐛 Describe the bug

nvJPEG leaks memory and fails with OOM after ~1-2k images.

import torch
from torchvision.io import read_file, decode_jpeg

for i in range(1000): # increase to your liking till gpu OOMs (:
    img_u8 = read_file('lena.jpg')
    img_nv = decode_jpeg(img_u8, device='cuda')

Probably related to first response to #3848

RuntimeError: nvjpegDecode failed: 5

is exactly the message you get after OOM.

Versions

PyTorch version: 1.9.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 11.1.0
Clang version: 12.0.1
CMake version: version 3.21.1
Libc version: glibc-2.33

Python version: 3.8.7 (default, Jan 19 2021, 18:48:37) [GCC 10.2.0] (64-bit runtime)
Python platform: Linux-5.13.8-arch1-1-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: 11.4.48
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 2080 Ti
GPU 1: NVIDIA GeForce RTX 2080 Ti
GPU 2: NVIDIA GeForce GTX 1080

Nvidia driver version: 470.57.02
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.2.2
/usr/lib/libcudnn_adv_infer.so.8.2.2
/usr/lib/libcudnn_adv_train.so.8.2.2
/usr/lib/libcudnn_cnn_infer.so.8.2.2
/usr/lib/libcudnn_cnn_train.so.8.2.2
/usr/lib/libcudnn_ops_infer.so.8.2.2
/usr/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] adabelief-pytorch==0.2.0
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.5
[pip3] pytorch-lightning==1.4.5
[pip3] torch==1.9.0+cu111
[pip3] torchaudio==0.9.0
[pip3] torchfile==0.1.0
[pip3] torchmetrics==0.4.1
[pip3] torchvision==0.10.0+cu111
[conda] Could not collect

[vadimkantorov]

@fmassa Is general torchvision.io.read_image affected? In general, is torchvision.io.read_image ready for being used in dataset classes instead of PIL / OpenCV for reading images from disk? Is it fast?

[Scass0807]

@NicolasHug @fmassa Also having this issue. Tried loading images on loop using decode_jpeg directly to GPU. Was able to run a number of images but GPU memory keeps going until it runs out and fails. Memory is fine on CPU. Was wondering if there is a timeline for when this will be fixed. Hoping it will be fixed ASAP as loading directly to GPU is crucial to getting speeds fast enough to run in real time.

[rydenisbak]

same problem:
ubuntu 20.04
Driver Version: 470.82.00
CUDA Version: 11.4
torch version '1.10.2+cu113'
torchvision version '0.11.3+cu113'

[NicolasHug]

Thanks all for the reports.

I took a look a this today. I can reproduce the leak. I do see the memory usage going up constantly with nvidia-smi. But printing the allocated and reserved cuda memory with torch.cuda.memory_stats() shows no leak: I always get

CUDA rsvrd = 12.583 MB
CUDA alloc = 12.583 MB

I thought the leak might come from the fact that we don't free the nvjpeg handle (we literally leak it for convenience)

vision/torchvision/csrc/io/image/cuda/decode_jpeg_cuda.cpp

Lines 28 to 30 in 9ae0169

	namespace {
	static nvjpegHandle_t nvjpeg_handle = nullptr;
	}

but that's not the case: putting back the handle within the function and properly destroying it with nvjpegDestroy() still leads to a leak.

I don't see the leak anymore when commenting out the nvjpegDecode() call, so it happens somewhere inside. I don't understand why nvjpegDecode() has to allocate anything though, because the output cuda tensor's memory was already allocated prior to the call.

I don't know whether that's actually a bug from nvjpeg, or if there's something else going on. Either way, I don't understand. nvjpeg allows to pass custom device memory allocators, perhaps there is something to do there.
I'll try to look into that, but meanwhile if anyone has any idea of what's going on, I would appreciate any help.

Cheers

[NicolasHug]

nvjpeg allows to pass custom device memory allocators, perhaps there is something to do there.
I'll try to look into that

Update: this still leaks 🥲

int dev_malloc(void **p, size_t s) {
    *p = c10::cuda::CUDACachingAllocator::raw_alloc(s);
    return 0;
  }

int dev_free(void *p) {
    c10::cuda::CUDACachingAllocator::raw_delete(p);
    return 0;
}

...

  nvjpegDevAllocator_t dev_allocator = {&dev_malloc, &dev_free};
  nvjpegStatus_t status = nvjpegCreateEx(NVJPEG_BACKEND_DEFAULT, &dev_allocator,
                                         NULL, NVJPEG_FLAGS_DEFAULT,  &nvjpeg_handle);

source for ref: https://github.com/NicolasHug/vision/blob/3f057677f5676bf12c42571f6bf72686dd9c61f2/torchvision/csrc/io/image/cuda/decode_jpeg_cuda.cpp#L28

[NicolasHug]

I had a chance to look at this more: this is an nvjpeg bug. Unfortunately I'm not sure we can do much about it.

It was fixed with CUDA 11.6 but I'm still observing the leak with 11.0 - 11.5.

A temporary fix for linux users is to download the 11.6 nvjpeg.so e.g. from here and to tell ld to use it instead of whatever you currently have installed (using LD_LIBRARY_PATH, or LD_PRELOAD, or something else)

[rydenisbak]

Hello @NicolasHug thanks for answer! I reinstalled CUDA, now I have this version
NVIDIA-SMI 510.47.03
Driver Version: 510.47.03
CUDA Version: 11.6
torch version '1.10.2+cu113'
torchvision version '0.11.3+cu113'

But problem does not disappear. Should I rebuild torch with cuda 11.6 from source?

[NicolasHug]

What does ldd <path_to_your_python_env>/site-packages/torchvision/image.so say regarding libnvjpeg? It's possible that it's still being linked to an 11.3 version, or that your .so file isn't as up to date as the I linked to above (11.6.0.55-1).

[Scass0807]

@NicolasHug Mine is showing /site-packages/torchvision/../torchvision.libs/libnvjpeg.90286a3c.so.11 How do I fix this to use system cuda?