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

self.visualization[name] = None
# Register the forward hook and keep the handle
# in another dict
self.handles[name] = \
layer.register_forward_hook(hook_fn)
def remove_hooks(self):
# Loops through all hooks and removes them
for handle in self.handles.values():
handle.remove()
# Clear the dict, as all hooks have been removed
self.handles = {}
setattr(StepByStep, 'attach_hooks', attach_hooks)
setattr(StepByStep, 'remove_hooks', remove_hooks)
The procedure is fairly straightforward now: Give it a list containing the names of
the layers to attach hooks to, and you’re done!
Hooking It
sbs_cnn1.attach_hooks(
layers_to_hook=['conv1', 'relu1', 'maxp1', 'flatten',
'fc1', 'relu2', 'fc2']
)
To get the visualization attribute filled with values, we still need to make
predictions:
Making Predictions (Logits)
images_batch, labels_batch = iter(val_loader).next()
logits = sbs_cnn1.predict(images_batch)
Don’t forget to remove the hooks after you’re finished with the predictions. By the
way, you can call predict() multiple times, and the outputs produced by the
hooked layers will be concatenated.
402 | Chapter 5: Convolutions