Interaction Nvidia DALI and Squirrel

[axkoenig]

Description

Describes in detail how Squirrel and DALI can work together. Also includes benchmarks on how to best utilize DALI and how it compares to transforms in Torchvision.

Attaching PDF rendered version of the Sphinx documentation here. Unfortunately, I couldn't get syntax highlighting to work.

Apparent next steps are figuring out how Squirrel and DALI can work together in multi-processing. It is not obvious how we could implement this, and if this provides a performance boost. Using a DALI parallel external source would probably be the way to go, but DALI expects a callable here that fetches individual images given a specific image index. This can be implemented easily if we set shard-size=1, but our initial experiments showed that larger shard sizes are more desirable.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring including code style reformatting
• Other (please describe):

Checklist:

• I have read the contributing guideline doc (external contributors only)
• Lint and unit tests pass locally with my changes
• I have kept the PR small so that it can be easily reviewed
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• All dependency changes have been reflected in the pip requirement files.

[github-actions]

This is PR is marked as stale as it has been inactive for 30 days. It will be closed in 7 days.

[pzdkn]

Nice work! LGTM!

Its surprising that increasing threads doesn't have much impact. Do you have an intuition on this ?

[axkoenig]

Hi @pzdkn,
thanks for taking a look at this. Squirrel already does multi-threading under the hood by default to fetch the shards from the disk, so I believe that's why adding even more threads with DALI doesn't help. What is very interesting, is the multi-processing case, where I would expect a significant boost in overall loading speed, but it is not trivial to implement, because DALI expects a source that fetches samples by data index (afaik you can think of it as a map-style dataset), which is not obvious to implement in squirrel (unless your shard size is 1, then you can simply fetch each key - but as we know having such small shard sizes will slow squirrel down, so unsure what the overall boost here would be).
Anyway, thanks again - merging now.
Cheers,
a