How to integrate with PyTorch DataLoaders

Aberta
  • Algo errado?
#23 3 comentários 0 reações 1 responsável Ver no GitHub

@jantonguirao já está trabalhando nisso.

Desde 29/10/2024.

Avaliação

Esta issue ainda não foi avaliada.

Descrição

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
Linguagem predominante
Jupyter Notebook
Estrelas
164
Forks
26
Merge médio
27min
PRs com merge (30d)
1

Preparar o ambiente

Primeiros passos

  1. Leia a issue inteira e depois o guia de contribuição do projeto.
  2. Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
  3. Faça um fork do repositório e trabalhe em uma branch.
  4. Abra um pull request que referencie o número da issue.

Mais de NVIDIA/nvImageCodec

Todas as issues de NVIDIA/nvImageCodec

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.