将子模块放到一个字典里
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.