Deep Learning with PyTorch

Extracting the losses

Just like we extracted the activation of a convolution layer using the register_forward_hook() function in Chapter 5, Deep Learning for Computer Vision, we can extract losses of different convolutional layers required to calculate style loss and content loss. The one difference in this case is that instead of extracting from one layer, we need to extract outputs of multiple layers. The following class integrates the required change:

class LayerActivations():    features=[]        def __init__(self,model,layer_nums):                self.hooks = []        for layer_num in layer_nums:            self.hooks.append(model[layer_num].register_forward_hook(self.hook_fn))        def hook_fn(self,module,input,output):        self.features.append(output)    def remove(self): for hook ...