Add default arg values to JAX decorator #5115
Conversation
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Signed-off-by: Albert Wolant <awolant@nvidia.com>
|
!build |
|
CI MESSAGE: [10415392]: BUILD STARTED |
|
CI MESSAGE: [10415392]: BUILD PASSED |
| if 'num_threads' not in wrapper_kwargs: | ||
| wrapper_kwargs['num_threads'] = 4 |
There was a problem hiding this comment.
This choice seems quite arbitrary. I'd recommend getting the number of processors (or perhaps the number of processors per GPU?).
There was a problem hiding this comment.
I did that because that's what we do it a lot in our codebase and especially in the tutorials. Do we have some benchmarks of what it does to the performance?
There was a problem hiding this comment.
But you need to account for the number of CPUs that may need to be used for the JAX itself, the speed of the network, and things that may run in parallel.
I'm afraid that without an elaborate autotune we can't provide an accurate value here.
There was a problem hiding this comment.
There's a lot of middle ground between elaborate autotune and a fixed number.
Still, I'd at least like to have a big, fat comment saying that it's subject to change.
There was a problem hiding this comment.
Done, added docs, comments and tutorial section on num_threads
Signed-off-by: Albert Wolant <awolant@nvidia.com>
|
!build |
|
CI MESSAGE: [10473251]: BUILD STARTED |
|
CI MESSAGE: [10473251]: BUILD PASSED |
Category:
New feature
Description:
Adds default values for
device_idandnum_threadsfor decorated JAX iterator functions. Thanks to this change it is easier to applyshardingto scale up the computation in the most common pattern.Additional information:
Affected modules and functionalities:
JAX iterator decorator.
Key points relevant for the review:
Tests:
Checklist
Documentation
DALI team only
Requirements
REQ IDs: N/A
JIRA TASK: DALI-3671