Is this the correct way to pass Haiku model for evaluation?

[bigyankarki]

I initialized the Haiku model as follows:

model, params, st = init_model(rng)

I do my gradient updates as follows:

params, opt_state, loss = update(params, opt_state, (batch_state, batch_y_policy, batch_y_value))

and send the model and params for evaluation
evaluate(model, params, st, rng)

and do a forward pass and calculate loss and accuracy

(predict_policy, predict_value), _st = model.apply(params=params, state=st, x=state[:batch_size].reshape(batch_size, 4, 4, 1), rng=rng)
policy_loss = jnp.mean(optax.softmax_cross_entropy(labels=batch[1], logits=y_predict_policy))
value_loss = jnp.mean(optax.l2_loss(y_value, y_predict_value))

For evaluation:

def evaluate(model, params, st, rng):
    # load test dataset
    test_dataset = np.load('test_set.npy', allow_pickle=True)
    print('Loaded test dataset with {} samples'.format(test_dataset.shape[1]))

    # batch the dataset
    batch_size = test_dataset.shape[1] # a single batch

    # convert from strucutred array to jax array
    state = test_dataset[0]
    y_policy = test_dataset[1]
    y_value = test_dataset[2]
    state = jnp.array([state[i] for i in range(len(state))], dtype=jnp.float32)
    y_policy = jnp.array([y_policy[i] for i in range(len(y_policy))], dtype=jnp.float32)
    y_value = jnp.array([y_value[i] for i in range(len(y_value))], dtype=jnp.float32)

    # evaluate model
    (predict_policy, predict_value), _st = model.apply(params=params, state=st, x=state[:batch_size].reshape(batch_size, 4, 4, 1), rng=rng)
    policy_loss = jnp.mean(optax.softmax_cross_entropy(labels=y_policy, logits=predict_policy))
    value_loss = jnp.mean(optax.l2_loss(y_value, predict_value))


    # calculate accuracy
    y_predict_value = predict_value.flatten() # flatten y_predict_label
    y_predict_value = jnp.where(y_predict_value <= 0, -1, 1)  # set y_predict_label to -1 if value is <=0  or else 1
    y_value = y_value.flatten() # flatten y_label

    # calculate accuracy
    accuracy = jnp.mean(jnp.where(y_value == y_predict_value, 1, 0))

and I just call evaluate(model, params, st, rng).

Is this a right way to pass model? Because it seems, although loss seems to be decreasing during training, it is stable during evaluation. So, I suspect that updated params are not being used.

[LenaMartens]

From a first look this seems to be correct, if evaluate gets updated parameters, it should be using them in the model.apply call. Not sure why you would get different losses. Is there a more minimal runnable example you could provide?

Since this is about Haiku and not JAX, could you open an issue on our issue tracker (https://github.com/deepmind/dm-haiku/issues) to follow up? Feel free to @ tag me in the issue and I can follow up there.

[bigyankarki]

I raised an issue in the haiku repo, and tagged you there. Thank you :)