Best Practise for Loading Data onto GPU with Jax?

[LawrenceMMStewart]

Hello everyone,

I hope you are all well. I have a question with respect to the optimal way of handling a dataset for a model trained on GPU.

Suppose I have a very large data set X which is too large to fit on the GPU, and in addition I have my model's parameters params sitting on the GPU as a jax.numpy array. More concretely, lets suppose X is a numpy array with shape = ( no_batches, batch_size, dimension of data) so X[i] is the i'th batch of training data, and that a training step consists of running a function train_step with the parameters and batch as input.

Currently I would move a batch onto the GPU and then initiate a training step on this batch, i.e.

for i in range(no_batches):
       # move batch onto the GPU by creating a jnp.array
       batch = jnp.array(X[i])
       train_step(params, batch)

However, this seems inefficient, as for each step in the loop the function, train_step must wait for the batch to be moved onto the GPU, when in fact the next batch could be moved onto the GPU in preparation whilst train_step is executing with the current batch.

What is the optimal way to go about moving data onto the GPU for training with JAX?

Many thanks for your help.

[cgarciae]

Two comments:

1. To my understanding, what you did is not inefficient since JAX works async, neither jnp.array nor your train_step will block unless you have a blocking operation like print, things just get queued for later execution in the background as they become available.
2. Use jnp.asarray, it sometimes prevents unnecessary copying.

[LawrenceMMStewart]

To my understanding, what you did is not inefficient since JAX works async, neither jnp.array nor your train_step will block unless you have a blocking operation like print, things just get queued for later execution in the background as they become available.

Thanks for the clear answer. Do you know where I can find a list / explanation as to what operations block (such as print)?

Use jnp.asarray, it sometimes prevents unnecessary copying.

Great, thats good to know.

Many thanks for your help @cgarciae !