如何在 Pytorch 中迭代层

Question

假设我有一个名为 m 的网络模型对象。现在我没有关于这个网络层数的先验信息。如何创建一个 for 循环来遍历其层？我正在寻找类似的东西：

 Weight=[]
for layer in m._modules:
    Weight.append(layer.weight)

原文由 Infintyyy 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

假设您有以下神经网络。

 import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 1 input image channel, 6 output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # define the forward function
        return x

现在，让我们打印与每个 NN 层关联的权重参数的大小。

 model = Net()
for name, param in model.named_parameters():
    print(name, param.size())

输出：

 conv1.weight torch.Size([6, 1, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])

我希望您可以扩展示例以满足您的需求。

原文由 Wasi Ahmad 发布，翻译遵循 CC BY-SA 4.0 许可协议