[HOW] stax log intermediate layer outputs

[llCurious]

I have tried flax before and there is a way to capture the intermediate outputs following the link https://flax.readthedocs.io/en/latest/howtos/extracting_intermediates.html#use-capture-intermediates.

Recently, i try to use stax which i think is more friendly to customized layers. And i wonder if there is a way that i could log the intermediate outputs in each layer (or some layer) in a similar way?

[tomhennigan]

Stax is not actively developed any more and does not have this functionality (plumbing intermediate values) built in.

There is a separate library available called Oryx which has a transformation called Harvest which can be used with any JAX code (e.g. you could use this with Stax, Flax, whatever). This transformation allows you to capture intermediate values and and collect them at the end of your function.

[tomhennigan]

You would need to insert a stax layer that had no parameters and called harvest on the input. This is actually what elementwise does so you can probably use that (I've not tested this):

nn_init, nn_apply = stax.serial(
        stax.Conv(2, (1, 1)),
        stax.elementwise(oryx.core.sow, name='layer_0', tag='intermediate'),
       ...
)

[llCurious]

Yep, the elementwise works fine.
However, it seems that the grad and reap cannot be done in the same time?

def loss_func(params, x, label):
    y = nn_apply(params, x)
    return mse_loss(y, label)

loss, gradients = jax.value_and_grad(loss_func)(opt_get_weights(opt_state), X, Y)
print('Reap:', oryx.core.reap(jax.value_and_grad(loss_func), tag='intermediate')(opt_get_weights(opt_state), X, Y))  # Pulls out 'layer_0'

[tomhennigan]

Seems like for value and reaped values you can use oryx.core.harvest(..) tensorflow/probability#1248 and some more discussion here tensorflow/probability#1045.

Perhaps @sharadmv has some thoughts?

[krzysztofrusek]

@tomhennigan is sow/harvest compatible with haiku transformations?
Ideally, I would like to sow intermediates inside hk.transform and harvest them otside.
Now I am using state for this, and oryx would be more elegant but I think it needs lifted sow aka hk.vmap etc.

PS

It works !
her is gist, though I am not sure if this is side effect free.

[llCurious]

thank you for your suggestion. The method you present works indeed, while the following approach is enough to sow and reap the intermediate values.

You would need to insert a stax layer that had no parameters and called harvest on the input. This is actually what elementwise does so you can probably use that (I've not tested this):

nn_init, nn_apply = stax.serial(
        stax.Conv(2, (1, 1)),
        stax.elementwise(oryx.core.sow, name='layer_0', tag='intermediate'),
       ...
)

What i am trying to do next is to obtain the intermediate, along with the the value and gradients which are more important. The current approach seems cannot obtain both simutaneously.