Monitoring training logs

[patrickvonplaten]

I've run into the problem of how to monitor values that are calculated within the loss function since the loss function is expected to only return a single loss values. E.g. imagine you have a trainining loop:

# some flax model `model`
model_weights = model.params

# define the train step
def train_step(model_weights, batch):
      model.params = model_weights
 
      # define the loss function
      def loss_fn(model_weights):
          labels = batch.pop("labels")
          
          # forward
          model_output = batch(batch)

          # get model outputs
          logits = model_outputs.logits
          
          # what if the model has an important metric that we would like to monitor, but not use as a loss 
          perplexity = model_output.perplexity

          # compute loss
          loss = jnp.linalg.norm(model(batch) - labels)
          return loss  # <- how can I return the perplexity here? `value_and_grad(...)` differentiates only by one value
      
      # compute grad function
      grad_fn = jax.value_and_grad(loss_fn)

      # apply grad function
      loss, grad = grad_fn(model_weights)
      grad = jax.lax.pmean(grad, "batch")

      apply_gradients(model_weights, grad)
      return model_weights, loss

# compile and map
p_train_step = jax.pmap(train, "batch")

...

# train
for i in range(num_train_steps):
    batch = data_collator(i)
    model_weights, loss = p_train_step(model_weights, batch)

=> Now how is it possible to log other metrics (perplexity) than just the loss? It's not that straight-forward in jax since the loss function is differentiated just from one output. Luckily jax.value_and_grad(...) has a special flag for this, called has_aux=True (see: https://jax.readthedocs.io/en/latest/jax.html#jax.grad) which allows us to monitor additional parameters

So the following above lines have to be changed:

return loss => return loss, perplexity
grad_fn = jax.value_and_grad(loss_fn) => grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
loss, grad = grad_fn(model_weights) => (loss, perplexity), grad = grad_fn(model_weights)

A more in-detail example can be found here:

flax/examples/lm1b/train.py

Line 220 in 9350b44

grad_fn = jax.value_and_grad(loss_fn, has_aux=True)

[andsteing]

Yes the has_aux=True solution is what we use in multiple examples (mostly for returning logits that are then used for more metric computations), e.g.

flax/examples/mnist/train.py

Line 62 in 8611538

grad_fn = jax.value_and_grad(loss_fn, has_aux=True)

flax/examples/imagenet/train.py

Line 142 in 8611538

aux, grads = grad_fn(state.params)

flax/examples/wmt/train.py

Line 226 in 8611538

grad_fn = jax.value_and_grad(loss_fn, has_aux=True)

Consider renaming the discussion to "Returning multiple values from loss function" or similar to the tip more easily discoverable (since it's not so much about monitoring but more about returning multiple values)