【Pytorch】获取网络的任意一层的输出

一、卷积和反卷积的计算

卷积层特征图计算

class torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, bias=True)

feature-map height :

feature-map width :

(stride[0]表示横向移动步长],stride[1]表示纵向移动步长])

反卷积层特征图计算

class torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, bias=True)  

二、获取网络的任意一层的输出

1.如果一个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].

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

3.通过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'])