如何从pytorch模块获取子模块序列?

How to obtain sequence of submodules from a pytorch module?

对于 pytorch module,我想我可以使用 .named_children.named_modules 等来获取子模块列表。但是,我想列表没有按顺序给出,对吧?一个例子:

In [19]: import transformers

In [20]: model = transformers.DistilBertForSequenceClassification.from_pretrained('distilb
    ...: ert-base-cased')

In [21]: [name for name, _ in model.named_children()]
Out[21]: ['distilbert', 'pre_classifier', 'classifier', 'dropout']

上述模型中.named_children()的顺序分别为distilbert、pre_classifier、classifier、dropout。但是,如果您检查 code,很明显 dropout 发生在 classifier 之前。那么如何获取这些子模块的顺序呢?

在 Pytorch 中,print(model).named_children() 等的结果是根据它们在模型的 class 的 __init__ 中声明的顺序列出的,例如

案例一

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
        self.conv2_drop = nn.Dropout2d()

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=0.6)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Model()
print(model)
[name for name, _ in model.named_children()]
# output
['conv1', 'conv2', 'fc1', 'fc2', 'conv2_drop']

案例二

更改了构造函数中 fc1fc2 层的顺序。

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc2 = nn.Linear(50, 10)
        self.fc1 = nn.Linear(320, 50)
        self.conv2_drop = nn.Dropout2d()

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=0.6)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Model()
print(model)
[name for name, _ in model.named_children()]
# output
['conv1', 'conv2', 'fc2', 'fc1', 'conv2_drop']

这就是为什么 classifierdropout 之前打印的原因,因为它在构造函数中是这样声明的:

class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
        ...
        self.distilbert = DistilBertModel(config)
        self.pre_classifier = nn.Linear(config.dim, config.dim)
        self.classifier = nn.Linear(config.dim, config.num_labels)
        self.dropout = nn.Dropout(config.seq_classif_dropout)

尽管如此,您可以使用 .modules() 等来使用模型的子模块,但它们只会按照它们在 __init__ 中声明的顺序列出。如果你只想打印基于 forward 方法的结构,你可以尝试使用 pytorch-summary.