Gathering activations that occur inside a neural net during training

Question

I'd like to use the activations of some network layer that occur during training for visualization and analysis of network performance. I had a look at both TrainingProgressFunction and TrainingProgressReporting but these do not give access to the inner workings of the net. Is it possible to get hold of activations while they occur e.g. using Sow and Reap?

I guess I am looking for something like Google's Tensorboard that lets one attach a node to some layer that collects the tensors during training (much like what Mathematica does with Sow and Reap during usual evaluation of an expression)

Answer 1

This seems like an interesting feature to add, given that it should be fairly easy to implement. What actual data would you be wanting about the interior activations? Statistical properties such as the RMS power of activations at a particular layer? Or the raw activations for you to process in your own way via a callback function? How expensive would this function be? Would you want this function to be applied at every single batch, or would you want to do it periodically, such as after each round, on a randomly selected batch? Or do you want to look at activations on some hold-out data that isn't involved in training?

The most general but most expensive thing is something we want to add anyway for the next version, which is to allow you to 'tap off' a particular interior output when applying a network via the second argument. Here's how it would work:

chain = NetInitialize @ NetChain[{2,Ramp,3,Ramp}, "Input" -> {}];
chain[5] (* usual output, e.g. output of final layer *)
chain[5, NetPort[2,"Output"]] (* output of 2nd layer *)

If this feature existed you could then accomplish what you want with:

interiorActivations[net_] := Module[{act},
    act = net[$mydata, NetPort[...]];
    dosomething[act];
];

NetTrain[..., TrainingProgressFunction -> Function[interiorActivations[#Net]]]

It's not as elegant but you can even do something equivalent in the current version. You can use NetGraph to make the interior activations you care about available as an output port. You won't train on that output port (so you'll have to specify a third argument to NetTrain if you aren't already), but you will have access to that port via #Net in a custom callback, as above.

Be aware, btw, that TrainingProgressReporting -> {f, "Interval" -> ...} can also be used to produce a dynamically-updating visualization. It's got some undocumented smarts about how it lays out the return value of the callback f, e.g. if its an association it typesets it in a table, if its a packed array it gets formatted in a compact grid, etc.