查看Pytorch网络的各层输出(feature map)、权重(weight)、偏置(bias)

下载地址

BatchNorm2d参数量

torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
# 卷积层中卷积核的数量C 
num_features – C from an expected input of size (N, C, H, W)

>>> import torch
>>> m = torch.nn.BatchNorm2d(100)
>>> m.weight.shape
torch.Size([100])
>>> m.numel()
AttributeError: 'BatchNorm2d' object has no attribute 'numel'
>>> m.weight.numel()
100
>>> m.parameters().numel()
Traceback (most recent call last):
  File "", line 1, in <module>
AttributeError: 'generator' object has no attribute 'numel'
>>> [p.numel() for p in m.parameters()]
[100, 100]

linear层

>>> import torch
>>> m1 = torch.nn.Linear(100,10)
# 参数数量= (输入神经元+1)*输出神经元
>>> m1.weight.shape
torch.Size([10, 100])
>>> m1.bias.shape
torch.Size([10])
>>> m1.bias.numel()
10
>>> m1.weight.numel()
1000
>>> m11 = list(m1.parameters())
>>> m11[0].shape
# weight
torch.Size([10, 100])
>>> m11[1].shape
# bias
torch.Size([10])

weight and bias

# Method 1 查看Parameters的方式多样化，直接访问即可
model = alexnet(pretrained=True).to(device)
conv1_weight = model.features[0].weight

# Method 2 
# 这种方式还适合你想自己参考一个预训练模型写一个网络，各层的参数不变，但网络结构上表述有所不同
# 这样你就可以把param迭代出来，赋给你的网络对应层，避免直接load不能匹配的问题！
for layer,param in model.state_dict().items(): # param is weight or bias(Tensor) 
	print layer,param

feature map

由于pytorch是动态网络，不存储计算数据，查看各层输出的特征图并不是很方便！分下面两种情况讨论：

1、你想查看的层是独立的,那么你在forward时用变量接收并返回即可！！

class Net(nn.Module):
    def __init__(self):
        self.conv1 = nn.Conv2d(1, 1, 3)
        self.conv2 = nn.Conv2d(1, 1, 3)
        self.conv3 = nn.Conv2d(1, 1, 3)

    def forward(self, x):
        out1 = F.relu(self.conv1(x))
        out2 = F.relu(self.conv2(out1))
        out3 = F.relu(self.conv3(out2))
        return out1, out2, out3

2、你的想看的层在nn.Sequential()顺序容器中，这个麻烦些，主要有以下几种思路：

# Method 1 巧用nn.Module.children()
# 在模型实例化之后，利用nn.Module.children()删除你查看的那层的后面层
import torch
import torch.nn as nn
from torchvision import models

model = models.alexnet(pretrained=True)

# remove last fully-connected layer
new_classifier = nn.Sequential(*list(model.classifier.children())[:-1])
model.classifier = new_classifier
# Third convolutional layer
new_features = nn.Sequential(*list(model.features.children())[:5])
model.features = new_features

​

# Method 2 巧用hook,推荐使用这种方式，不用改变原有模型
# torch.nn.Module.register_forward_hook(hook)
# hook(module, input, output) -> None

model = models.alexnet(pretrained=True)
# 定义
def hook (module,input,output):
    print output.size()
# 注册
handle = model.features[0].register_forward_hook(hook)
# 删除句柄
handle.remove()

# torch.nn.Module.register_backward_hook(hook)
# hook(module, grad_input, grad_output) -> Tensor or None
model = alexnet(pretrained=True).to(device)
outputs = []
def hook (module,input,output):
    outputs.append(output)
    print len(outputs)

handle = model.features[0].register_backward_hook(hook)

注：还可以通过定义一个提取特征的类，甚至是重构成各层独立相同模型将问题转化成第一种

计算模型参数数量

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)