Incorrect hook being used in register_hook implementation

username: Toshi
cooked: <p>Hi, I’ve been trying to implement masking on gradients over some custom layer weights in my network, and I implement it in the following way (before start of each epoch/training step).</p>
<pre><code class="lang-auto">cur_hooks = {}    
for n,p in net.named_children():
    if isinstance(p, Conv2d) or isinstance(p, Linear):
            gradient_mask = (p.some_func(p.novel_para).round().float()&lt;=1e-6).data
            gradient_mask = gradient_mask.float()
            cur_hooks[n] = p.weight.register_hook(lambda grad: grad.mul_(gradient_mask))
</code></pre>
<p>If I print the hooks via</p>
<pre><code class="lang-auto">        for k,v in cur_hooks.items():
            print(k, v)
</code></pre>
<p>It prints:</p>
<pre><code class="lang-auto">conv1 &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92a1c0&gt;
conv1a &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92afa0&gt;
conv1b &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92a370&gt;
conv1c &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92a9a0&gt;
conv1d &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92aee0&gt;
conv2a &lt;torch.utils.hooks.RemovableHandle object at 0x7f554e05bd90&gt;
conv2b &lt;torch.utils.hooks.RemovableHandle object at 0x7f554e05bb80&gt;
shortcut_conv2 &lt;torch.utils.hooks.RemovableHandle object at 0x7f554e05be20&gt;
conv2c &lt;torch.utils.hooks.RemovableHandle object at 0x7f554e05bc70&gt;
conv2d &lt;torch.utils.hooks.RemovableHandle object at 0x7f554e05bd30&gt;
conv3a &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92f250&gt;
conv3b &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92f100&gt;
shortcut_conv3 &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92f9d0&gt;
conv3c &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92f430&gt;
conv3d &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92f880&gt;
conv4a &lt;torch.utils.hooks.RemovableHandle object at 0x7f554c92f2e0&gt;
conv4b &lt;torch.utils.hooks.RemovableHandle object at 0x7f554df74d90&gt;
shortcut_conv4 &lt;torch.utils.hooks.RemovableHandle object at 0x7f554df74f70&gt;
conv4c &lt;torch.utils.hooks.RemovableHandle object at 0x7f554df74cd0&gt;
conv4d &lt;torch.utils.hooks.RemovableHandle object at 0x7f554df74430&gt;
linear &lt;torch.utils.hooks.RemovableHandle object at 0x7f554df74f10&gt;
</code></pre>
<p>However, at the very first backward() call, it runs into error-</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/mnt/train.py", line 144, in &lt;module&gt;
    loss.backward()
  File "/home/user/base/lib/python3.9/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/user/base/lib/python3.9/site-packages/torch/autograd/__init__.py", line 130, in backward
    Variable._execution_engine.run_backward(
  File "/mnt/train.py", line 107, in &lt;lambda&gt;
    cur_hooks[n] = p.weight.register_hook(lambda grad: grad.mul_(gradient_mask))
RuntimeError: The size of tensor a (3) must match the size of tensor b (512) at non-singleton dimension 3
</code></pre>
<p>Apparently (If I’m not wrong) - its using the hook for the Linear layer (the last hook) for the first Convolution layer (conv1) - and I can’t seem to figure out what I did wrong. Any suggestions?</p>

username: ptrblck
cooked: <p>Are you getting the same shape mismatch error if you are using <code>torch.ones_like(grad)</code> instead of <code>gradient_mask</code>?<br>
If not, then I assume the <code>gradient_mask</code> creation is creating the wrong shape.<br>
If yes, could you post a minimal, executable code snippet, which would reproduce this error?</p>

username: Toshi
cooked: <p>Here is a dummy script to recreate the error.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision.datasets as dset
import torchvision.transforms as transforms

from torchvision.transforms import Compose, ToTensor, RandomAffine, RandomApply, Resize
import torch.nn.functional as F

import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5x5 image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = LeNet().to(device=device)

## train, test eval:
loss_metric = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), 0.01)

dataset = dset.FashionMNIST(root='/mnt/sda1/Data/', download=True, train=True,
                            transform=transforms.Compose([
                                transforms.Resize(28),
                                transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,)),
                            ]))
valset = dset.FashionMNIST(root='/mnt/sda1/Data/', download=True, train=False,
                            transform=transforms.Compose([
                                transforms.Resize(28),
                                transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,)),
                            ]))
train_loader = torch.utils.data.DataLoader(dataset, batch_size=128,
                                                shuffle=True, num_workers=4)

# testloader = torch.utils.data.DataLoader(valset, batch_size=128,
#                                                 shuffle=False, num_workers=4)
def freeze_weights(model, hooks = []):
    for name, module in model.named_children():
        if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
            gradient_mask = (module.weight&gt;.05).data
            gradient_mask = gradient_mask.float()
            h_ = module.weight.register_hook(lambda grad:print\
            ("Pre hook shape {}: non zero % {}, gradient mask shape {}".format\
            (grad.shape, torch.count_nonzero(grad)/torch.numel(grad), gradient_mask.shape)))
            h = module.weight.register_hook(lambda grad: grad.mul_(gradient_mask))
            h__ = module.weight.register_hook(lambda grad:print\
            ("Post hook shape {}: non zero % {}".format\
            (grad.shape, torch.count_nonzero(grad)/torch.numel(grad))))
            hooks = hooks + [h] + [h_] + [h__]
            print("{} gradient mask shape {} when created".format(name, gradient_mask.shape))
    if len(hooks)==0:
        raise Exception("hook list is empty!") # should be 15
    return hooks
    
for epoch in range(10):
    if epoch&gt;0:
        curr_hooks = freeze_weights(net)
        print("{} hooks created".format(len(curr_hooks)))
    
    for i, (x,y) in enumerate(train_loader):
        optimizer.zero_grad()

        x, y = x.to(device), y.to(device)
            
        outputs = net(x)

        loss = loss_metric(outputs,y)
        
        loss.backward()
        optimizer.step()
    
    if epoch&gt;0:
        for h in curr_hooks: h.remove()
        curr_hooks = []

</code></pre>
<p>This is on Python 3.9.1 and torch version 1.7.1+cu110.</p>
<p>From what I understood, the gradient mask of the last layer/hook is being used in all of the hooks. I think I’d have to define a dict() for the gradient masks to keep track of them?</p>

username: Toshi
cooked: <p>I think the solution is</p>
<pre><code class="lang-auto">def freeze_weights(model, hooks = [], gradient_mask_weights ={}):
    ...
    ...
    h = module.weight.register_hook(lambda grad, n=name: grad.mul_(gradient_mask_dict[n]))
    ...
    ...
    return hooks, gradient_mask_weights
</code></pre>

username: ptrblck
cooked: <p>Thanks for the code snippet.<br>
The approach using a <code>dict</code> might work, alternatively, you could also use:</p>
<pre><code class="lang-python">h = module.weight.register_hook(lambda grad, gradient_mask=gradient_mask: grad.mul_(gradient_mask))
</code></pre>
<p>Your current code runs into a known Python limitation of using <code>lambda</code> functions, since names in function bodies are evaluated when the function is executed.</p>

username: Mriganka_Nath
cooked: <p>is there any difference between register_hook and register_backward_hook? both of them says registers a backward hook.<br>
Whats the differnce</p>

username: ptrblck
cooked: <p><a href="https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html"><code>register_hook</code></a> is used on a Tensor, while <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=register_backward_hook#torch.nn.Module.register_backward_hook"><code>register_backward_hook</code></a> is used on <code>nn.Module</code>, is deprecated, and should be replaced with <code>register_full_backward_hook</code>.</p>

username: winterknight
cooked: <p>I want to use activations an intermediate layer in the T5 decoder to compute a separate loss. Is there a way of adding a custom layer on an intermediate layer of a pretrained model and using the output of both the new layer head as well as the pretrained model head ?</p>
<p>I tried getting the activations of an intermediate layer of a pretrained model using hooks. But running into errors (sharing below).</p>
<p>More specifically I am unable to grab the value of activations. For eg, in the following code snippet, <code>summarizer_decoder_activation</code> is empty.</p>
<pre><code class="lang-auto">model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict=True)
summarizer_decoder_activation = {}
def get_activation(name):
  def hook(model, input, output):
      summarizer_decoder_activation[name] = output.detach()
  return hook

for name, layer in model._modules["decoder"]._modules["block"][-1]._modules["layer"][-1]._modules["DenseReluDense"].named_children():
# for name, layer in model.named_children():
  if name == "wo":
    layer.register_forward_hook(get_activation(name))
</code></pre>

username: ptrblck
cooked: <p>You are detaching the intermediate activation, so it’ll be a constant in the loss calculation.<br>
Remove it and it should work:</p>
<pre><code class="lang-python">activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output
    return hook


model = models.resnet18()

model.layer1[1].bn2.register_forward_hook(get_activation("aux_out"))

x = torch.randn(2, 3, 224, 224)
out = model(x)
print(out.shape)
# torch.Size([2, 1000])
print(activation["aux_out"].shape)
# torch.Size([2, 64, 56, 56])

loss = out.mean() + activation["aux_out"].mean()
loss.backward()
</code></pre>