Неизвестное поведение HOOKS в PyTorch - PullRequest
Новая
       7

Неизвестное поведение HOOKS в PyTorch

0 голосов
/ 28 мая 2020

У меня есть простой и понятный 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)

Эти крючки были созданы с использованием следующих 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))

Обратите внимание, что между Hook 1 и Hook 2 нет ReLU self.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()

1 Ответ

1 голос
/ 28 мая 2020

Проблема в том, что в вашем Conv2d слое input - это кортеж, а output - torch.Tensor. Следовательно, output[0][0,...] выбирает первый элемент из dim 0 в тензоре, тогда как input[0][0,...] выбирает первый элемент из кортежа.

Вам просто нужно изменить output[0][0,...] на output[0,...].

...