Unknown behaviour of HOOKS in PyTorch

Question

I have a straightforward and simple CNN below,

# creat a dummy deep net
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(1,2, kernel_size=3, stride=1, padding=1, bias=True)
        self.conv2 = nn.Conv2d(2,3, kernel_size=3, stride=1, padding=1, bias=True)
        self.conv3 = nn.Conv2d(3,1, kernel_size=3, stride=1, padding=1, bias=True)
        self.seq = nn.Sequential(
                    nn.Conv2d(1,5, kernel_size=3, stride=1, padding=1, bias=True),
                    nn.LeakyReLU(negative_slope=0.2, inplace=True),
                    nn.Conv2d(5,1, kernel_size=3, stride=1, padding=1, bias=True),
                    )
        self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):

        out = self.relu(self.conv1(x))
        out = self.conv3(self.conv2(out))
        out = out + x
        out = self.seq(x)

        return out

5 hooks have been applied to each layer for the forward pass.

Hooked 0 to Conv2d(1, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Hooked 1 to Conv2d(2, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Hooked 2 to Conv2d(3, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Hooked 3 to Sequential(
  (0): Conv2d(1, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): LeakyReLU(negative_slope=0.2, inplace=True)
  (2): Conv2d(5, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)
Hooked 4 to LeakyReLU(negative_slope=0.2, inplace=True)

These hooks have been created using following class

# ------------------The Hook class begins to calculate each layer stats
class Hook():
    def __init__(self, module, backward=False):
        if backward==False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)

        self.inputMean = []
        self.outputMean = []

    def hook_fn(self, module, input, output):
        self.inputMean.append(input[0][0,...].mean().item())#calculate only for 1st image in the batch
        print('\nIn hook class input {}'.format(input[0].size()))
        self.outputMean.append(output[0][0,...].mean().item())
        print('In hook class outout {}'.format(output[0].size()))

# create hooks on each layer
hookF = []
for i,layer in enumerate(list(net.children())):
    print('Hooked to {}'.format(layer))
    hookF.append(Hook(layer))

Please note between Hook 1 and Hook 2 there is no ReLU self.conv3(self.conv2(out)). Thus OUTPUT of HOOK1 is INPUT to HOOK2 and should be identical. BUT THIS DOES NOT TURNS OUT TO BE WHY? Below is output for HOOK1 and HOOK2

Hook of layer 1 (HOOK on layer 1 which is self.conv2)
... OutputMean: [0.2381615787744522, 0.2710852324962616, 0.30706286430358887, 0.26064932346343994, 0.24395985901355743]

 Hook of layer 2 (HOOK on layer 2 which is self.conv3)
InputMean: [0.13127394020557404, 0.1611362248659134, 0.1457807868719101, 0.17380955815315247, 0.1537724733352661], OutputMean: ...

These two values should have been the same but do not turn out to be.

------ The Full code is shown below -------

import torch
import torch.nn as nn

# creat a dummy deep net
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(1,2, kernel_size=3, stride=1, padding=1, bias=True)
        self.conv2 = nn.Conv2d(2,3, kernel_size=3, stride=1, padding=1, bias=True)
        self.conv3 = nn.Conv2d(3,1, kernel_size=3, stride=1, padding=1, bias=True)
        self.seq = nn.Sequential(
                    nn.Conv2d(1,5, kernel_size=3, stride=1, padding=1, bias=True),
                    nn.LeakyReLU(negative_slope=0.2, inplace=True),
                    nn.Conv2d(5,1, kernel_size=3, stride=1, padding=1, bias=True),
                    )
        self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):

        out = self.relu(self.conv1(x))
        out = self.conv3(self.conv2(out))
        out = out + x
        out = self.seq(x)

        return out

net = Net()
print(net)
criterion = nn.MSELoss()


# ------------------The Hook class begins to calculate each layer stats
class Hook():
    def __init__(self, module, backward=False):
        if backward==False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)

        self.inputMean = []
        self.outputMean = []

    def hook_fn(self, module, input, output):
        self.inputMean.append(input[0][0,...].mean().item())#calculate only for 1st image in the batch
        print('\nIn hook class input {}'.format(input[0].size()))
        self.outputMean.append(output[0][0,...].mean().item())
        print('In hook class outout {}'.format(output[0].size()))

# create hooks on each layer
hookF = []
for i,layer in enumerate(list(net.children())):
    print('Hooked to {}'.format(layer))
    hookF.append(Hook(layer))

optimizer = torch.optim.Adam(net.parameters())

# Do 5 forward pass
for _ in range(5):
    print('Iteration --------')
    data = torch.rand(2,1,10,10)*10
    print('Input mean is {}'.format(data[0,...].mean()))
    target = data.clone()

    out = net(data)
    loss = criterion(out, target)
    print('backward')
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

for i,h in enumerate(hookF):    
    print('\n Hook of layer {}'.format(i))
    print('InputMean: {}, OutputMean: {}'.format(h.inputMean, h.outputMean))
    h.hook.remove()

Answer 1

The problem is that in your Conv2d layer input is a tuple and output is a torch.Tensor. Therefore output[0][0,...] is selecting the first item from dim 0 in the tensor whereas input[0][0,...] is selecting the first item from the tuple.

You just need to change output[0][0,...] to output[0,...].