Pytorch学习（十六）----获取网络的任意一层的输出

其实一般来说，如果一个net中，是一个Sequential直接包起来，首先直接print（net ）即可，然后看到类似：

    (net1): Sequential(
      (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))
      (1): ReLU()
      (2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))
      (3): ReLU()
      (4): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))
      (5): ReLU()
      (6): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))
      (7): ReLU()
      (8): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))
      (9): ReLU()
      (10): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2))
      (11): ReLU()
      (12): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))
    )

其中 net1是属性，在定义net的时候，net1用sequential定义了。我们想要拿到net1的第三层层直接 net.net1[2], 因为net1是一个list, 因此这里用[2].

另外一种思路就是，用过个sequential来弄：

比如：

import torch
from torch.autograd import Variable
import torch.nn

class my_net(nn.Module):
    def __init__(self):
        super(my_net, self).__init__()
        self.features1 = torch.nn.Sequential(
                            torch.nn.Conv2d(3, 5, 3, padding=1),
                            torch.nn.ReLU(),
                            torch.nn.Conv2d(5, 10, 3, padding=1),
                            )
        # Take the output of this layer as the input of net2
        self.features2 = torch.nn.Sequential(
                            torch.nn.ReLU(),
                            torch.nn.Conv2d(10, 15, 3, padding=1)
                            )
    def forward(self, x):
        x1 = self.features1(x)
        x2 = self.features2(x1)

        # 呃呃呃呃呃， 对，通过这种方式进行返回，虽然ugly，但是straightforward
        return (x1,x2)

class my_net2(nn.Module):
    def __init__(self):
        super(my_net2, self).__init__()
        self.features = torch.nn.Sequential(
                            torch.nn.Conv2d(10, 25, 3, padding=1)
                            )
    def forward(self, x):
        x = self.features(x)
        return x

net1 = my_net().cuda()
net2 = my_net2().cuda()


input1 = Variable(torch.ones(1, 3, 10, 10).cuda())
feat_for_net2, out1 = net1(input1)

# input of net2 are features got from net1
input2 = feat_for_net2
out2 = net2(input2)

# Create targets
target1 = Variable(torch.ones_like(out1.data))
target2 = Variable(torch.ones(1, 25, 10, 10)*2)


criterion1 = torch.nn.MSELoss(size_average=False)
criterion2 = torch.nn.MSELoss(size_average=False)

loss1 = criterion1(out1, target1)

loss2 = criterion2(out2, target2)

loss = loss1+loss2
loss.backward()

另外一种是通过hook的方式，

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.cl1 = nn.Linear(25, 60)
        self.cl2 = nn.Linear(60, 16)
        self.fc1 = nn.Linear(16, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        
    def forward(self, x):
        x = F.relu(self.cl1(x))
        x = F.relu(self.cl2(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.log_softmax(self.fc3(x), dim=1)
        return x


activation = {}
def get_activation(name):
    def hook(model, input, output):
        # 如果你想feature的梯度能反向传播，那么去掉 detach（）
        activation[name] = output.detach()
    return hook


model = MyModel()
model.fc2.register_forward_hook(get_activation('fc2'))
x = torch.randn(1, 25)
output = model(x)
print(activation['fc2'])

网络定义时的不同方法：

1. 如果是单纯一个大的sequential，里面的每一层的名子其实会自动编号的。就是不直观，不过写forward时候比较容易些，直接类似：self.net1(x)就行了。
2. 如果是将net1拆成每层来写，那么forward就比较拖沓，不过直观。

最后，如果是网络在多GPU运行的那种，就是被 DataParallel包了一下的。你直接 net.module就行了，就可以拿到真正的网络。