skorch doctor: a tool to understand the net

[BenjaminBossan]

I recently watched Andrej Karpathy's new series (yes, I still feel like a noob :)) and also remembered his blog post from a few years ago. These sources contain a couple of useful tips for better understanding the neural net training and diagnose potential problems, shining some light into the blackbox.

I wondered if we could not automate some of these steps to make it super easy for users to attain those insights. So I started working on automating those steps that make sense to be automated. The results can be seen in the notebook here.

This is obviously still in a very rough shape, but before spending more time on it, I wanted to get feedback from @ottonemo and @thomasjpfan if they think it's useful to have something like this.

If we want to proceed, I would implement this in a proper module with tests and docs, and also test it with a more complex architecture (BERT or something along those lines).

For now, the new allows to train with some data and automatically records some useful data. As a user, I can just do:

doctor = SkorchDoctor(net)
doctor.fit(X_sample, y_sample)
doctor.plot_loss()
doctor.plot_activations()
doctor.plot_gradients()
# etc.

I also dabbled into creating something similar to sklearn's cv_resutls_, i.e. a report with some summary statistics. Not sure how useful that would be.

Since SkorchDoctor adds a bunch of hooks and modifies the net in a few ways. This will create a memory overhead, but I try to pull as much to down from GPU to numpy as possible. It's also better to create a new net instance afterwards, though I try to clean up the hooks and modifications as best as possible.

Ah yes, I also found a small bug in on_grad_computed, which has training=False as default arg but obviously during training, it should be set to True. It has no consequences in our code since it's only used by GradientNormClipping, which doesn't check the argument, but still...

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[BenjaminBossan]

@ottonemo @thomasjpfan This feature is now ready for review. Please find below the full description, which can also be used for the squashed commit message.

Description

A helper class to assist in understanding the neural net training

The SkorchDoctor helper class allows users to wrap their neural net before
training and then automatically collect useful data that allows to better
understand what is going on during training and how to possibly improve it.

The class will automatically record activations of each module, gradients, and
updates of each learnable parameter, all of those for each training step. Once
training is finished, the user can either directly take a look at the data,
which is stored as an attribute on the helper class, or use one of the provided
plotting functions (requires matplotlib) to plot distributions of the data.

Examples of what conclusions could be drawn from the data:

• Net is not powerful enough
• Need for better weight initialization or normalization
• Need to adjust optimizer
• Need for gradient clipping

However, the helper class will not suggest any of those solutions itself, I
don't think that's possible. It is only intended to help surfacing potential
problems, it's up to the user to decide on a solution.

A notebook to show the usage of SkorchDoctor, once for a simple MLP and once for
fine-tuning a BERT model, is provided:

https://github.com/skorch-dev/skorch/blob/skorch-doctor/notebooks/Skorch_Doctor.ipynb

Implementation

Because of the additional data being collected, depending on the use case, a
significant memory overhead is expected. To keep this in check, a few measures
are taken:

• The collected data is immediately pulled to numpy to avoid clobbering GPU
  memory.
• It is documented, and shown in examples, that you should use only a small
  amount of data and low number of epochs, since that's enough to understand
  most problems. Most notably, this helps with storing less data about
  activations.
• For parameter updates, only a single scalar per weight/bias is stored,
  indicating the relative magnitude of the update.
• The biggest overhead will most likely come from storing the gradients, not
  sure if something can be done here without losing too much useful data.

For storing activations, some heuristics are in place to deal with the output of
the modules. The problem here is that modules can return any arbitrary data from
their forward call. A few assumptions are being made here: The output can be
shoved into to_numpy and it has to be either a torch tensor, a list, a tuple, or
a mapping of torch tensors. If it's neither of those, an error is raised.

Coincidental changes

While working on this PR, I stumbled upon a few minor bugs that are fixed as
part of the PR:

1. Inferring the predict nonlinearity made the assumption that a net.criterion_
   exists. However, this might not always be the case. Now, this function works
   even if the criterion attribute has a different name. Moreover, when more
   than one criterion is defined, the identity function will be returned.
2. The second bug is that in check_is_fitted, we made the check dependent on the
   module_ attribute. Again, we should not assume that it always exists, as
   users may define different names. Now, those custom names will be checked,
   and only if those don't exist is it assumed that the module_ attribute should
   exist.
3. The helper.py module now defines an __all__ attribute. This seems to be
   necessary for sphinx to build documentation for objects that are imported to,
   but not defined in, helper.py. Specifically, this is the case for the new
   SkorchDoctor class, but also the existing parse_args CLI feature.
4. Fixed a small bug in on_grad_computed, which has training=False as default
   arg but during training, it should be set to True. Thankfully, it had no
   consequences in our code since it's only used by GradientNormClipping, which
   doesn't check that argument.

[thomasjpfan]

Can the coincidental changes be split out into their own PRs so they can be merged with independent of this PR?

As for including skorch doctor, are there existing libraries that hook onto a module and record the same type of information to visualize?

[BenjaminBossan]

Can the coincidental changes be split out into their own PRs so they can be merged with independent of this PR?

I did: #927

As for including skorch doctor, are there existing libraries that hook onto a module and record the same type of information to visualize?

I am not aware of any library that would provide similar functionality. As recording all this extra information is a memory overhead, no library would provide this by default. It only makes sense to log this with the specific purpose of performing the type of analysis provided here.

The sources that inspired me to add the feature are the ones linked, i.e. the video series/blog post by Karpathy, the borealis.ai blog post, and the experience that colleagues and I gathered when trying to diagnose training issues.

[thomasjpfan]

I recently got a chance to watch the videos. It was very interesting!

This is a first pass over the code itself.

[thomasjpfan]

Can the logs be renamed "records"? I keep thinking that "log_grad" means the mathematical "log(grad)".

[thomasjpfan]

Do you think it's useful to have two kwargs to pass into the train and val loss separately?

[BenjaminBossan]

Hmm, then we couldn't just pass them as **kwargs but would require them to be passed as dicts, which would result in a different API from all the other plotting functions. I see the benefit but I'm not sure if they outweigh the costs. Since the ax is returned to the user at the end, they can still do some changes to the plot afterwards.

[thomasjpfan]

kwargs is routed directly to the hist plot. What do you think of the following API?

def plot_activation(...,
    hist_kwargs=None,
):
    hist_kwargs = hist_kwargs or {}

This comment applies to the other plot_* methods.

[BenjaminBossan]

Not sure if I quite understand, you want to avoid having **kwargs at all but rather prefer to have foo_kwargs where foo corresponds to the kind of plot being generated? One reason why I chose to implement kwargs like this is because it is similar to the pandas plotting API, where you can also pass arbitrary kwargs and they are passed to the underlying plotting function.

[thomasjpfan]

I may have missed it, but what was the reason behind log10 compared to natural log (ln)?

[BenjaminBossan]

I don't think there is any strict reason, it's just easier to translate log10 in your head I guess. When you look at the corresponding part of the video, you can see that Andrej even starts with log and then changes to log10 :)

[thomasjpfan]

Should the axes be flattened for cases where the user passes in axes that is "2d":

fig, axes = plt.subplots(3, 4)

doctor.plot_activations(..., axes=axes)

[thomasjpfan]

I'm trying to find where SkorchDoctor adds a callback. May you point to the code where this happens?

[BenjaminBossan]

Ah yes, an earlier implementation added an extra callback, but after a refactoring, this was no longer necessary. Therefore, the test can now be removed.

[BenjaminBossan]

Thanks for the review Thomas, I made some changes according to your comments or asked some clarifying questions.

Regarding the name, good point about the potentially confusing name of "logs". I changed it now to "recs". I didn't choose "records" because it's longer, I hope that works for you.

[BenjaminBossan]

Hmm, then we couldn't just pass them as **kwargs but would require them to be passed as dicts, which would result in a different API from all the other plotting functions. I see the benefit but I'm not sure if they outweigh the costs. Since the ax is returned to the user at the end, they can still do some changes to the plot afterwards.

[BenjaminBossan]

Not sure if I quite understand, you want to avoid having **kwargs at all but rather prefer to have foo_kwargs where foo corresponds to the kind of plot being generated? One reason why I chose to implement kwargs like this is because it is similar to the pandas plotting API, where you can also pass arbitrary kwargs and they are passed to the underlying plotting function.

[BenjaminBossan]

I don't think there is any strict reason, it's just easier to translate log10 in your head I guess. When you look at the corresponding part of the video, you can see that Andrej even starts with log and then changes to log10 :)

[BenjaminBossan]

Ah yes, an earlier implementation added an extra callback, but after a refactoring, this was no longer necessary. Therefore, the test can now be removed.

[thomasjpfan]

The update with match_fn makes sense. I left one small testing commit, otherwise I'm okay with merging.

[thomasjpfan]

I know that an error would raise if nothing matched, but can we also assert len(doctor.activation_recs_["module"]) to make sure there is something in the list?

[BenjaminBossan]

Done