Python PyTorch ModuleDict用法及代码示例

本文简要介绍python语言中 torch.nn.ModuleDict 的用法。

用法:

class torch.nn.ModuleDict(modules=None)

参数：

modules(可迭代的,可选的) -(字符串：模块)的映射(字典)或类型(字符串，模块)的键值对的可迭代

在字典中保存子模块。

ModuleDict 可以像常规 Python 字典一样进行索引，但它包含的模块已正确注册，并且所有 Module 方法都可见。

ModuleDict是一个排序尊重的字典

• 插入顺序，以及

• 在 update() 中，合并 OrderedDict 、 dict (从 Python 3.6 开始)或另一个 ModuleDict (update() 的参数)的顺序。

请注意，update() 与其他无序映射类型(例如，Python 3.6 版本之前的 Python 普通 dict)不会保留合并映射的顺序。

例子：

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.choices = nn.ModuleDict({
                'conv': nn.Conv2d(10, 10, 3),
                'pool': nn.MaxPool2d(3)
        })
        self.activations = nn.ModuleDict([
                ['lrelu', nn.LeakyReLU()],
                ['prelu', nn.PReLU()]
        ])

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x