Daniel Voigt Godoy - Deep Learning with PyTorch Step-by-Step A Beginner’s Guide-leanpub

IMPORTANT: The logging function must return None, otherwise
the gradients will be modified, assuming the returned value.
The register_hook() method registers a backward hook to a
tensor for a given parameter. The hook function takes a
gradient as input and returns either a modified gradient or None.
The hook function will be called every time a gradient with
respect to that tensor is computed.
Since we’re using this function for logging purposes, we should
leave the gradients alone and return None.
"Isn’t there a register_backward_hook() method? Why aren’t we
using it?"
That’s a fair question. At the time of writing, this method still has an unsolved issue,
so we’re following the recommendation of using register_hook() for individual
tensors instead.
Now, we can use the new method to log gradients for the linear layer of our model,
never forgetting to remove the hooks after training:
Model Training
1 sbs_adam = StepByStep(model, loss_fn, optimizer)
2 sbs_adam.set_loaders(train_loader)
3 sbs_adam.capture_gradients('linear')
4 sbs_adam.train(10)
5 sbs_adam.remove_hooks()
By the time training is finished, we’ll have collected two series of 50 gradients each
(each epoch has five mini-batches), each series corresponding to a parameter of
linear (weight and bias), both of them stored in the _gradients attribute of our
StepByStep instance.
Learning Rates | 465