CUDA Error/RuntimeError when using decode_jpeg() in dataloader

Coprosmo: <p>I’m trying to decode a large number of jpeg images directly onto my GPU via a DataLoader, but am running into issues.</p>
<p>I’m running with <code>batch_size=64</code> and <code>num_workers=8</code>.</p>
<p>The code runs as expected when I pass <code>device="cpu"</code> to <code>decode_jpeg()</code>.</p>
<p>My dataset is defined as follows:</p>
<pre><code class="lang-auto">class MyDataset(Dataset):
    def __init__(self):
        super(MyDataset).__init__()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.files = sorted(glob.glob("my_dir/*.jpeg"))

    def __getitem__(self, idx):
        im_u8 = read_file(self.files[idx])
        im_nv = decode_jpeg(im_u8, device=self.device).float() / 255
        return im_nv

    def __len__(self):
        return len(self.files)
</code></pre>
<p>Initially I was running into a CUDA initialization error, which I fixed by setting <code>mp.set_start_method("spawn")</code>. I am now running into the following cryptic RuntimeError:</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “/code/scripts/predict-noncapture-hauls.py”, line 161, in <br>
run()<br>
File “/usr/local/lib/python3.9/dist-packages/click/core.py”, line 1130, in <strong>call</strong><br>
return self.main(*args, **kwargs)<br>
File “/usr/local/lib/python3.9/dist-packages/click/core.py”, line 1055, in main<br>
rv = self.invoke(ctx)<br>
File “/usr/local/lib/python3.9/dist-packages/click/core.py”, line 1404, in invoke<br>
return ctx.invoke(self.callback, **ctx.params)<br>
File “/usr/local/lib/python3.9/dist-packages/click/core.py”, line 760, in invoke<br>
return __callback(*args, **kwargs)<br>
File “/code/scripts/predict-noncapture-hauls.py”, line 147, in run<br>
predictions, metadata = evaluate_event_yolov5(<br>
File “/code/scripts/predict-noncapture-hauls.py”, line 75, in evaluate_event_yolov5<br>
for batch in dataloader:<br>
File “/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py”, line 530, in <strong>next</strong><br>
data = self._next_data()<br>
File “/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py”, line 1224, in _next_data<br>
return self._process_data(data)<br>
File “/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py”, line 1250, in _process_data<br>
data.reraise()<br>
File “/usr/local/lib/python3.9/dist-packages/torch/_utils.py”, line 457, in reraise<br>
raise exception<br>
RuntimeError: Caught RuntimeError in DataLoader worker process 0.<br>
Original Traceback (most recent call last):<br>
File “/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/worker.py”, line 287, in _worker_loop<br>
data = fetcher.fetch(index)<br>
File “/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/fetch.py”, line 52, in fetch<br>
return self.collate_fn(data)<br>
File “/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/collate.py”, line 136, in default_collate<br>
storage = elem.storage()._new_shared(numel)<br>
File “/usr/local/lib/python3.9/dist-packages/torch/_tensor.py”, line 181, in storage<br>
storage = self._storage()<br>
RuntimeError: it != attype_to_py_storage_type.end()INTERNAL ASSERT FAILED at “…/torch/csrc/DynamicTypes.cpp”:69, please report a bug to PyTorch. Failed to get the Python type of <code>_UntypedStorage</code>.</p>
</blockquote>

ptrblck: <p>What type is <code>im_nv</code> using? It seems as if the type cannot be mapped to a storage.</p>

Coprosmo: <p><code>im_nv</code> is of type <code>torch.float32</code> and shape 3x640x640. <code>im_u8</code> has type <code>torch.uint8</code>.</p>
<p>Note that it works as expected when decoding to CPU.</p>

Coprosmo: <p><strong>Edit</strong>: Replied earlier saying I was no longer receiving this issue. It turns out I’d increase the batch size, which was causing a different issue. I am still receiving the RuntimeError.</p>
<p>This is really odd. I put a flag in after the <code>decode_jpeg()</code> line, which showed that it is actually successfully loading the first few images (possibly all of the images in each worker’s batch?), and then breaking.</p>
<p>It looks like it’s breaking at the collate step - could it be an issue with the tensors in each batch being on gpu prior to collation?</p>

Coprosmo: <p>Update: moving the tensors back to CPU in <code>__getitem__</code> stops the error from occurring. Though this isn’t ideal since we then move the batches back onto GPU.</p>

Coprosmo: <p>A further strange observation: When I load the jpeg in via OpenCV, turn it into a tensor, and move that onto GPU, it works fine as expected.</p>
<p>For reference, loading an image via <code>torchvision.io.decode_jpeg()</code>:</p>
<pre><code class="lang-auto">im_u8 = read_file(self.files[idx])
im_nv = decode_jpeg(im_u8, device="cuda").float() / 255
</code></pre>
<p>versus via OpenCV:</p>
<pre><code class="lang-auto">im = cv2.imread(self.files[idx])  # BGR
im = im.transpose((2, 0, 1))[::-1]
im = np.ascontiguousarray(im)
im = torch.from_numpy(im).float() / 255
im = im.to("cuda")
</code></pre>
<p>The produced tensors are not exactly the same (minor decoding differences I guess) but are close. They are the same shape and data type. The OpenCV option works as expected, while the PyTorch one produces our friend the RuntimeError <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>
<p>(Note, I tried making <code>im_nv</code> contiguous to no avail)</p>

ptrblck: <p>Thanks for the follow-up! I’m currently unsure what’s causing the issue, but could you create an issue on the <code>torchvision</code> <a href="https://github.com/pytorch/vision/issues">repository</a> so that we can track and debug the issue further, please?</p>

Coprosmo: <p>Sure, will condense it down a bit <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20">. Thanks for your help so far!</p>

ptrblck: <p>I tried to reproduce the issue using <code>torchvision==0.13.0+cu116</code> (the RC) with a single image and this code:</p>
<pre><code class="lang-python">import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
import torch.multiprocessing as mp

class MyDataset(Dataset):
    def __init__(self):
        super(MyDataset).__init__()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.path = 'fake.jpeg'

    def __getitem__(self, idx):
        im_u8 = torchvision.io.read_file(self.path)
        im_nv = torchvision.io.decode_jpeg(im_u8, device=self.device).float() / 255
        return im_nv

    def __len__(self):
        return 10
    
def main():
    dataset = MyDataset()
    loader = DataLoader(dataset, num_workers=2)
    
    for data in loader:
        print(data)
        
if __name__=="__main__":
    mp.set_start_method("spawn")
    main()
</code></pre>
<p>which seems to work fine. Could you also try out the latest nightly release?</p>

Coprosmo: <p>I just tried your example with the nightly release, and am no longer getting the RuntimeError. Unfortunately I’m running into the memory leak bug known to occur before CUDA 11.6 (although my machine does seem to be running 11.6 <img src="https://discuss.pytorch.org/images/emoji/apple/person_shrugging.png?v=12" title=":person_shrugging:" class="emoji" alt=":person_shrugging:" loading="lazy" width="20" height="20">).</p>
<p>I wasn’t able to find the distribution of torchvision you mentioned. I can only see the CUDA 10.2 and 11.3 versions mentioned on the website, could you let me know how you installed it?</p>

ptrblck: <p>This should work:</p>
<pre><code class="lang-python">pip install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cu116
</code></pre>
<p>Could you also point me to the known memory leak issue in previous versions, please?</p>

Coprosmo: <p>Here is a a thread on the memory leak issue: <a href="https://github.com/pytorch/vision/issues/4378" class="inline-onebox" rel="noopener nofollow ugc">Torchvision decode_jpeg memory leak · Issue #4378 · pytorch/vision · GitHub</a></p>
<p>I’m no longer getting the warning on <code>torchvision==0.13.0+cu116</code>, but am still running into memory issues. Using a smaller batch size stops the GPU memory issues, though this is substantially slower than running the CPU decoding with a larger batch size, for some reason.</p>
<p>I am also receiving the following warning when decoding on GPU:</p>
<blockquote>
<p>[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]</p>
</blockquote>
<p>I’m currently investigating to see whether this could be related to the memory leak issues.</p>