可以只获取 PyTorch 模型中父组件的名称

It is possible to get only the names of the parent components in PyTorch model

Pytorch 中的所有预训练模型都包含具有预定义名称的“父”子模块,例如 AlexNet 包含 3 个“父”子模块:featuresavgpoolclassifier

model = torch.hub.load('pytorch/vision:v0.10.0','resnet101',pretrained=True)
AlexNet(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): ReLU(inplace=True)
    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (9): ReLU(inplace=True)
    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(6, 6))
  (classifier): Sequential(
    (0): Dropout(p=0.5, inplace=False)
    (1): Linear(in_features=9216, out_features=4096, bias=True)
    (2): ReLU(inplace=True)
    (3): Dropout(p=0.5, inplace=False)
    (4): Linear(in_features=4096, out_features=4096, bias=True)
    (5): ReLU(inplace=True)
    (6): Linear(in_features=4096, out_features=1, bias=True)
  )
)

是否有一种方法可以只获取这些组件,例如 mo = model.get_subs() # mo=['features','avgpool','classifier']

你可以使用这个:

import torch
import torchvision.models as models

model = models.alexnet(pretrained=True)
parents = [parent[0] for parent in model.named_children()] # get parents names
print(parents)

输出:

['features', 'avgpool', 'classifier']