ModuleDict¶
- class torch.nn.ModuleDict(modules=None)[source]¶
在一个字典中保存子模块。
ModuleDict
可以像普通的 Python 字典一样索引,但它包含的模块将被正确注册,并且所有Module
方法都可以看到。ModuleDict
是一个有序字典,它尊重插入顺序,以及
在
update()
中,合并的OrderedDict
、dict
(从 Python 3.6 开始)或另一个ModuleDict
(update()
的参数)的顺序。
请注意,使用其他无序映射类型(例如,在 Python 3.6 版本之前的 Python 的普通
dict
)的update()
不会保留合并映射的顺序。- 参数
modules (iterable, optional) – (字符串: 模块) 的映射(字典)或类型为 (字符串, 模块) 的键值对的迭代器
示例
class MyModule(nn.Module): def __init__(self) -> None: super().__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
- update(modules)[source]¶
使用映射中的键值对更新
ModuleDict
,覆盖现有键。注意
如果
modules
是一个OrderedDict
、一个ModuleDict
或一个键值对的可迭代对象,则其中新元素的顺序将被保留。