Pytorch获取神经网络各层的参数名称及参数值

一、parameters()

model.parameters()可以返回model网络的可学习参数，具体代码如下：

import torch.nn as nn

# 定义网络
class Net(nn.Module):
  def __init__(self):
    # nn.Module的子类必须在构造函数中执行父类的构造函数
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 6, 5) # 依次为输入通道、输出通道和卷积核
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(400, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

net = Net() # 创建网络
# print(net) # 打印网络

# 仅查看网络中的可学习参数（遍历）
for params in net.parameters():
    # print(params)
    # print(params.data) # 获取纯数据
    print(params.shape)

也可以通过下述方式进行获取：

import torch.nn as nn

class Net(nn.Module):
  def __init__(self):
    # nn.Module的子类必须在构造函数中执行父类的构造函数
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 6, 5) # 依次为输入通道、输出通道和卷积核
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(400, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

net = Net() # 创建网络
# print(net) # 打印网络

# 仅查看网络中的可学习参数，list形式
params = list(net.parameters())
# print(params) # 各层的参数值
# print(params[-2])
# print(params[-2].data) # 获取纯数据（最后一个全连接层的权重）
print(params[-2].data.shape) # 最后一个全连接层的权重的形状

二、named_parameters()

model.named_parameters()可同时返回model的可学习的参数及名称，具体代码如下：

import torch.nn as nn

# 定义网络
class Net(nn.Module):
  def __init__(self):
    # nn.Module的子类必须在构造函数中执行父类的构造函数
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 6, 5) # 依次为输入通道、输出通道和卷积核
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(400, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

net = Net() # 创建网络
# print(net) # 打印网络

# 查看网络中的可学习参数及名称
for name,params in net.named_parameters():
    print(name)
    # print(params)
    # print(params.shape)

三、state_dict()

model.state_dict()常见于模型的保存，它返回的是一个字典，分别对应模型中的键和值。

import torch.nn as nn

class Net(nn.Module):
  def __init__(self):
    # nn.Module的子类必须在构造函数中执行父类的构造函数
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 6, 5) # 依次为输入通道、输出通道和卷积核
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(400, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

net = Net() # 创建网络
# print(net) # 打印网络

# state_dict()存放训练过程中需要学习的权重和偏置
for k,v in net.state_dict().items(): # items()返回可遍历的(键, 值) 元组数组
    print(k)
    # print(v) # 纯数值
    # print(v.shape)

四、children()

model.children()可以返回网络的各层（最外层）。

import torch.nn as nn

class Net(nn.Module):
  def __init__(self):
    # nn.Module的子类必须在构造函数中执行父类的构造函数
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 6, 5) # 依次为输入通道、输出通道和卷积核
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(400, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

net = Net() # 创建网络
# print(net) # 打印网络

# net.children()遍历net的层（最wai层）
for x in net.children():
    print(x)

五、一些总结

在复现顶会论文开源代码时，经常遇到上述几个函数，因此特将它们记录下来。为了便于结果展示，我注释了很多代码，建议大家把这些注释去掉，都实践一下，多多体会。