-
Notifications
You must be signed in to change notification settings - Fork 80
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
What is unit of batch-size
and how should I set its value?
#426
Comments
The unit is the number of elements of There's no option to specify the batch size as a fraction of the input and we're currently looking into having a more robust implementation of a chunked-vmap (see google/jax#11319 or netket's implementation). |
As dataset size grows (either in sample count or feature dimension), the dataset itself can quickly become limited by GPU memory as well. In this case, it may be worth also exploring the use of a jax callback to bring (a chunk of) data to the device. |
Hi!
Following up on the solutions kindly provided, I'm experimenting with the
batch-size
method. It is very helpful! Now we are able to compute OT problems at least 10 times larger than before (still trying on larger sizes).batch-size
?It seems a bit counterintuitive. For example, if I'm computing the optimal transport between two distributions X_1 and X_2, I found that if I fix X_1 and increase the size of X_2 I would need to reduce the size of batch_size.
Right now, it seems I would need to do a manual line search for this parameter to find an "optimal feasible solution" for it :) Is there a more straightforward way to set its value?
Thankss!
The text was updated successfully, but these errors were encountered: