Support for a neural compression decoder

I recently started reading into neural/learned compression and implemented a simple CNN training pipeline using image data that I previously encoded/compressed with a neural compression model. Neural compression can achieve better compression factors than "classic" codecs and therefore be interesting for DNN training if, for example, main memory is limited and reading data from disk during training is too slow.

In my small test, I extended the PyTorch Dataset class and loaded a pretrained model from CompressAI into my Dataset to return images to the data loader. Unfortunately, in my pipeline the decoding time of the neural codecs I used were quite bad compared to JPEG and alike.

Therefore, I am wondering how neural decompression in DNN training pipelines can be made faster. One of my thoughts was if it would be possible to add support for a neural compression decoder op to DALI. From a user facing perspective I would envision an operator like nvidia.dali.fn.decoders.neural_codec() that takes a path to a trained decoder model checkpoint, maybe some type hints, and other common decoder args. The op would then instantiate the model and convert data passed from a reader op into tensors for downstream preprocessing.

I'd be curious to hear how feasible you see it to integrate a neural decoder into DALI. From a high-level perspective, I would consider it very helpful for anyone without unlimited main memory if more efficient neural compression methods were available as a drop-in replacement for classic codecs. Potential issues that I see based on my (very limited) knowledge of DALI internals are:

• It creates an additional dependency to a library that provides neural network primitives.
• DALI is agnostic of the downstream deep learning framework and I don't know if this would work if a user has to provide a trained model checkpoint to the neural decoder op.
• I'm not sure if GPUs could be efficiently used for this op since inference (decoding with a trained model) is often performed on CPUs.

Looking forward to your opinion!

No files, tests, or entry points are named. Start by reviewing DALI's operator architecture and the proposed PyTorch/CompressAI model-loading boundary, then assess the dependency and CPU/GPU execution concerns raised in the issue. Done would require an agreed implementation design and clear integration requirements.