ModuleDict¶
- class torch.nn.ModuleDict(modules=None)[来源][来源]¶
将子模块保存在一个字典中。
ModuleDict
可以像常规 Python 字典一样进行索引,但它包含的模块已正确注册,并且对所有Module
方法可见。ModuleDict
是一个有序字典,遵循插入顺序,以及
在
update()
中,合并的OrderedDict
、dict
(从 Python 3.6 开始)或另一个ModuleDict
(update()
的参数)的顺序。
请注意,使用其他无序映射类型(例如,Python 3.6 版本之前的普通
dict
)进行update()
操作不会保留合并映射的顺序。- 参数
modules (iterable, optional) – 一个(string: module)映射(字典)或一个类型为(string, module)的键值对可迭代对象
示例
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)[来源][来源]¶
使用来自映射的键值对更新
ModuleDict
,覆盖现有键。注意
如果
modules
是一个OrderedDict
、一个ModuleDict
或一个键值对可迭代对象,则保留其中新元素的顺序。