ModuleDict — PyTorch 2.7 documentation (original) (raw)
class torch.nn.ModuleDict(modules=None)[source][source]¶
Holds submodules in a dictionary.
ModuleDict can be indexed like a regular Python dictionary, but modules it contains are properly registered, and will be visible by allModule methods.
ModuleDict is an ordered dictionary that respects
- the order of insertion, and
- in update(), the order of the merged
OrderedDict
,dict
(started from Python 3.6) or anotherModuleDict (the argument toupdate()).
Note that update() with other unordered mapping types (e.g., Python’s plain dict
before Python version 3.6) does not preserve the order of the merged mapping.
Parameters
modules (iterable , optional) – a mapping (dictionary) of (string: module) or an iterable of key-value pairs of type (string, module)
Example:
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
Remove all items from the ModuleDict.
Return an iterable of the ModuleDict key/value pairs.
Return type
Iterable[tuple[str, torch.nn.modules.module.Module]]
Return an iterable of the ModuleDict keys.
Return type
Remove key from the ModuleDict and return its module.
Parameters
key (str) – key to pop from the ModuleDict
Return type
update(modules)[source][source]¶
Update the ModuleDict with key-value pairs from a mapping, overwriting existing keys.
Note
If modules is an OrderedDict
, a ModuleDict, or an iterable of key-value pairs, the order of new elements in it is preserved.
Parameters
modules (iterable) – a mapping (dictionary) from string to Module, or an iterable of key-value pairs of type (string, Module)
Return an iterable of the ModuleDict values.
Return type