将子模块放到一个字典里

torch.nn.ModuleDict(modules=None)

ModuleDict 可以像regular Python dictionary一样进行索引
ModuleDict 是一个有序的字典.

Note that update() with other unordered mapping types (e.g., Python’s plain dict) does not preserve the order of the merged mapping.

参数:

modules (iterable, optional) – a mapping (dictionary) of (string: module) or an iterable of key-value pairs of type (string, module)
作者注:
这个参数可以是字典、子模块、或者可迭代的键值对

示例

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

这个类中的方法:
1· clear()

clear()
清除ModuleDict中的所有item

2· items()

items()
返回一个可迭代的对象,里面包含一些ModuleDict键值对
Return an iterable of the ModuleDict key/value pairs.

3· keys()

keys()
返回ModuleDict中的键

4· pop(key)

pop(key)
Remove key from the ModuleDict and return its module.
参数:
key (string) – key to pop from the ModuleDict

5· update(modules)
在这里插入图片描述
6· values()

values()
Return an iterable of the ModuleDict values.
Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐