pytorch
4b2d1bec - [dynamo] keep submodule's name for nn.Sequential when unroolling (#94913)

Commit
2 years ago
[dynamo] keep submodule's name for nn.Sequential when unroolling (#94913) Currently, when unrolling an nn.Sequential, we use an integer to represent its submodule's name. This produces some difficulty in tracking the origin of the parameters in the export path: ```python model = nn.Sequential(OrderedDict([ ('conv1', nn.Conv2d(1,20,5)), ('relu1', nn.ReLU()), ('conv2', nn.Conv2d(20,64,5)), ('relu2', nn.ReLU()) ])) ``` Currently, the submodules will have names such as model.0, model.1 instead of model.conv1, model.relu1. This discrepency causes it difficult to track the origin of paramers because they are represented as model.conv1.foo and model.relu1.foo in model.named_parameters(). We replace enumerate() with named_children() to keep submodule's name. Pull Request resolved: https://github.com/pytorch/pytorch/pull/94913 Approved by: https://github.com/jansel
Author
Committer
Parents
Loading