The One PyTorch Trick Which You Should Know

How hooks can improve your workflow significantly

If you have ever used deep learning before, you know that debugging a model can be really hard sometimes. Tensor shape mismatches, exploding gradients, and countless other issues can surprise you. Solving these require looking at the model under the microscope. The most basic methods include littering the forward() methods with print statements or introducing breakpoints. These are of course not very scalable, because they require guessing where things went wrong, and are quite tedious to do overall.

However, there is a solution: hooks. These are specific functions, able to be attached to every layer and called each time the layer is used. They basically allow you to freeze the execution of the forward or backward pass at a specific module and process its inputs and outputs.

Let’s see them in action!

Hooks crash course

So, a hook is just a callable object with a predefined signature, which can be registered to any nn.Module object. When the trigger method is used on the module (i.e. forward() or backward() ), the module itself with its inputs and possible outputs are passed to the hook, executing before the computation proceeds to the next module.

In PyTorch, you can register a hook as a

• forward prehook (executing before the forward pass),
• forward hook (executing after the forward pass),
• backward hook (executing after the backward pass).

It might sound complicated at first, so let’s take a look at a concrete example!

An example: saving the outputs of each convolutional layer

Suppose that we want to inspect the output of each convolutional layer in a ResNet34 architecture. This task is perfectly suitable for hooks. In the next part, I will show you how can this be performed. If you would like to follow it interactively, you can find the accompanying Jupyter notebook at https://github.com/cosmic-cortex/pytorch-hooks-tutorial .

Our model is defined by the following.

Creating a hook to save outputs is very simple, a basic callable object is perfectly enough for our purposes.

An instance of SaveOutput will simply record the output tensor of the forward pass and stores it in a list.

A forward hook can be registered with the register_forward_hook(hook) method. (For the other types of hooks, we have register_backward_hook and register_forward_pre_hook .) The return value of these methods is the hook handle, which can be used to remove the hook from the module.

Now we register the hook to each convolutional layer.

When this is done, the hook will be called after each forward pass of each convolutional layer. To test it out, we are going to use the following image.

The forward pass:

As expected, the outputs were stored properly.

>>> len(save_output.outputs)
36

By inspecting the tensors in this list, we can visualize what the network sees.

Just for curiosity, we can check what happens later. If we go deeper in the network, the learned features become more and more high level. For instance, there is a filter which seems to be responsible for detecting the eyes.

Going beyond

Of course, this is just the tip of the iceberg. Hooks can do much more than simply store outputs of intermediate layers. For instance, neural network pruning, which is a technique to reduce the number of parameters, can also be performed with hooks .

To summarize, applying hooks is a very useful technique to learn if you want to enhance your workflow. With this under your belt, you’ll be able to do much more and do them more effectively.