How to integrate with PyTorch DataLoaders

Offen
  • Stimmt etwas nicht?
#23 3 Kommentare 0 Reaktionen 1 zugewiesene Person Auf GitHub ansehen

@jantonguirao arbeitet bereits daran.

Seit 29.10.2024.

Bewertung

Dieses Issue wurde noch nicht bewertet.

Beschreibung

question
Describe the question.

Hi,

I'm hoping to integrate nvImageCode with PyTorch DataLoaders (torch utils DataLoader, or FFCV DataLoader, or LitData DataLoader), but I'm struggling.

If I include the decoder as a transform to be used in my dataset.__getitem__ method, I get the dreaded cudaErrorInitializationError:
RuntimeError: Unhandled CUDA error: cudaErrorInitializationError initialization error

class CustomDataSet(StreamingDataset):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.decoder = nvimgcodec.Decoder() # cudaErrorInitializationError
    def __getitem__(self, idx):
        sample  = super().__getitem__(idx)
        sample['image'] = self.decoder(sample['image'])
        return sample

I can have my dataset return the raw image bytes, and apply the decoder to the list of bytes which is fast, but then I have to loop over items to transform them into pytorch tensors which is slow because it operates over the entire batch sequentially (not in parallel workers). This single step is slow enough that it negates the advantage of using the nvimgcodec.Decoder().

class CustomDataSet(StreamingDataset):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    def __getitem__(self, idx):
        sample  = super().__getitem__(idx)
        return sample

dataset = CustomDataSet(...)
dataloader = DataLoader(dataset, ...)

for batch in tqdm(dataloader):
    imgs = decoder.decode(batch['image'])
    imgs = [torch.tensor(img).moveaxis(-1,0) for img in imgs] # need to do this in the worker processes

I also tried to have my dataset return decode sources with ROIs, but this fails because DecodeSource is not pickleable.

class CustomDataset(StreamingDataset):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __getitem__(self, idx):
        sample  = super().__getitem__(idx) # <- Whatever you returned from the DatasetOptimizer prepare_item method.
        roi = nvimgcodec.Region(0, 0, 224, 224) # replace with random crop later...
        sample['image'] = nvimgcodec.DecodeSource(sample['image'], roi)
        return sample

In any case, I've checked open bugs/issues, and the docs, and I can't find a good example of using nvimgcodec in the context of a dataloader with parallel workers. Any guidance or suggestions for how to handle this would be greatly appreciated.

Check for duplicates
  • I have searched the open bugs/issues and have found no duplicates for this bug report
Vorherrschende Sprache
Jupyter Notebook
Sterne
164
Forks
26
Ø Merge
27 Min.
Gemergte PRs (30 T.)
1

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus NVIDIA/nvImageCodec

Alle Issues in NVIDIA/nvImageCodec

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.