Good afertnoon,
I just started learning JAX and Haiku but I can't run my code on GPU. I ran my code in Google Colab and Kaggle notebooks with GPU activated but it takes more time than when the GPU is deactivated.
Moreover when I look into the GPU metrics I see that I'm only using 1% of the computation power but 90% of the GPU memory.
Here is my code (MLP for MNIST) :
def mlp(images):
model = hk.Sequential([hk.Linear(128),
jax.nn.relu,
hk.Linear(64),
jax.nn.relu,
hk.Linear(10),
jax.nn.log_softmax])
return model(images)
def loss(params, model, images, labels):
logits = model.apply(params = params, images = images)
labels = jax.nn.one_hot(labels, num_classes = 10)
cross_entropy_loss = -jnp.sum(labels*logits)/len(labels)
return cross_entropy_loss
# Initializing the MLP model
mlp = hk.without_apply_rng(hk.transform(mlp))
params = mlp.init(rng = jax.random.PRNGKey(0),
images = next(iter(train_loader))[0])
# Initializing the optimizer
opt = optax.adam(1e-4)
opt_state = opt.init(params = params)
@jax.jit
def update(params, opt_state, images, labels):
grads = jax.grad(loss)(params,mlp,images,labels)
updates, opt_state = opt.update(grads, opt_state)
return optax.apply_updates(params, updates), opt_state
def train(params, opt_state, epochs):
for epoch in range(epochs):
for batch_idx, (images, labels) in enumerate(train_loader):
if batch_idx == 0:
print(f"Epoch {epoch} : loss = {loss(params,mlp,images,labels)}")
params, opt_state = update(params, opt_state, images,labels)
%time train(params, opt_state, epochs = 10)
If you have any idea of what I'm doing wrong you would help me a lot. Thank you.