pytorch查看模型每一层的名称

选与握

于 2024-07-30 19:10:49 发布

阅读量1.2k

点赞数 6

CC 4.0 BY-SA版权

分类专栏： # pytorch 文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/xihuanniNI/article/details/140803197

1 model._modules.items()

model._modules.items()是一个包含模型所有子模块的迭代器。在PyTorch中，当我们定义一个nn.Module子类时，我们可以使用nn.Sequential或nn.ModuleDict等容器类将多个子模块组合成一个整体。在这种情况下，我们可以通过访问nn.Module类中的_modules属性来访问这些子模块。_modules是一个有序字典，其中键是子模块的名称，值是子模块对象

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        return x

model = MyModel()
for name, module in model._modules.items():
    print(name, module)

输出如下:

conv1 Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
conv2 Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
relu ReLU()

在上面的例子中，我们可以通过迭代model._modules.items()来访问模型中的每个子模块。name是子模块的名称（例如"conv1"、“conv2"和"relu”），module是子模块的实例对象（例如nn.Conv2d和nn.ReLU）。在递归模型中，每个子模块都可以是一个递归模型，因此我们可以使用_modules属性来递归遍历整个模型的所有子模块。

2 named_modules()

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):