PyTorch中HOOKS的未知行为

Question

我下面有一个简单明了的CNN，

# 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个钩子以进行向前传递。

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)

这些钩子是使用下面的类创建的

# ------------------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))

请注意在挂钩1和挂钩2之间没有ReLUself.conv3(self.conv2(out))。因此，HOOK1的输出是HOOK2的输入，并且应该相同。但是这不是为什么吗？下面是HOOK1和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: ...

这两个值应该是相同的，但事实并非如此。

------完整代码如下-------

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

问题是，在您的Conv2d层中，input是一个元组，output是一个火炬张量。因此，output[0][0,...]从张量中的暗0中选择第一项，而input[0][0,...]从元组中选择第一项。

您只需要将output[0][0,...]更改为output[0,...]。