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

their gradient values.
For each monitored layer, it will go over its parameters, and, for those that require
gradients, it will create a logging function (log_fn()) and register a hook for it in
the tensor corresponding to the parameter.
The logging function simply appends the gradients to a list in the dictionary entry
corresponding to the layer and parameter names. The dictionary itself, _gradients,
is an attribute of the class (which will be created inside the constructor method, but
we’re setting it manually using setattr for now). The code looks like this:
StepByStep Method
setattr(StepByStep, '_gradients', {})
def capture_gradients(self, layers_to_hook):
if not isinstance(layers_to_hook, list):
layers_to_hook = [layers_to_hook]
modules = list(self.model.named_modules())
self._gradients = {}
def make_log_fn(name, parm_id):
def log_fn(grad):
self._gradients[name][parm_id].append(grad.tolist())
return None
return log_fn
for name, layer in self.model.named_modules():
if name in layers_to_hook:
self._gradients.update({name: {}})
for parm_id, p in layer.named_parameters():
if p.requires_grad:
self._gradients[name].update({parm_id: []})
log_fn = make_log_fn(name, parm_id)
self.handles[f'{name}.{parm_id}.grad'] = \
p.register_hook(log_fn)
return
setattr(StepByStep, 'capture_gradients', capture_gradients)
464 | Chapter 6: Rock, Paper, Scissors