Saving nn.Module to Parent nn.Module without Registering Paremeters

username: rambo
cooked: <p>Hi,</p>
<p>I have a parent nn.Module <code>parent</code> that saves multiple custom defined child nn.Modules as attributes of which I want one to not appear in <code>parent.children()</code>.</p>
<p>In other words, I want to save an nn.Module as an attribute of another nn.Module without registering its parameters to the former nn.Module. How can I do this?</p>
<p>Thanks!</p>

username: ptrblck
cooked: <p>If you are using a custom class (or just a Python <code>object</code>) for the children, I would guess that they won’t be registered in the parent <code>nn.Module</code>.</p>

username: rambo
cooked: <p>Hi, thanks for the answer!</p>
<p>As the modules arise systematically (and the relevant attribute is expected to be of type <code>Module</code> elsewhere), it would be nice to not have to wrap the child Module.</p>
<p>I had hoped there was some possibility to enter <code>__setattr__</code> with the child Module considered to be an object of type <code>object</code>, or would this still make the parent Module register the child in <code>children()</code>?</p>
<p>Thanks again!</p>

username: ptrblck
cooked: <p>I don’t think there is a clean approach to still use an <code>nn.Module</code>, but disallow registering it inside the parent. You could check <a href="https://github.com/pytorch/pytorch/blob/873255c6d95342d144e9d1b633c16410844b934e/torch/nn/modules/module.py#L366-L388">this code</a> and try to see, if changing the <code>self._modules</code> attribute might work, but note that this is an internal attribute and I don’t know what unwanted side effects you might trigger.<br>
Maybe you could explain your use case a bit more so that we could think about another approach?</p>

username: moi90
cooked: <p>I have the same requirement. My use case is that I want to add two versions of the same submodule, one with full functionality (metrics, additional computations for validation) and one stripped down for training only. Both share the same parameters:</p>
<pre><code class="lang-python">class MyModule(nn.Module):
  def __init__(self, m: nn.Module):
    self.m = m
    self._m_train = m.strip() # Remove any functionality not required for training

  def forward(self, input):
    return self._m_train(input)

my_module = MyModule(m)
optimizer = torch.optim.SGD(my_module.parameters())
</code></pre>
<p>Fortunately, <code>my_module.parameters()</code> doesn’t seem to contain duplicate parameters. But it seems cleaner (at least to me) to not have both versions of the same <code>m</code> registered as children.</p>

username: nishi951
cooked: <p>One quick and dirty hack I did was to wrap the submodule in tuple or list. For example:</p>
<pre><code class="lang-auto">class MyModule(nn.Module):
  def __init__(self, m: nn.Module):
    self._m = [m]
    self._m_train = m.strip() # Remove any functionality not required for training

  def forward(self, input):
    return self._m_train(input)
  
  # Optional
  @property
  def m(self):
    return self._m[0]

my_module = MyModule(m)
optimizer = torch.optim.SGD(my_module.parameters())
</code></pre>
<p>Btw this is nice when if you have a module that registers an altered version of itself as an attribute but avoid infinite recursion on module-recursing functions like <code>.to()</code> (for example).</p>

username: sammlapp
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="132082">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>could check <a href="https://github.com/pytorch/pytorch/blob/873255c6d95342d144e9d1b633c16410844b934e/torch/nn/modules/module.py#L366-L388" rel="noopener nofollow ugc">this code </a> and try to see, if changing the <code>self._modules</code> attribute might work, but note that this is an internal attribute and I don’t know what unwanted side effects you might trigger.<br>
Maybe you could explain your use case a bit more so that we could think about another approach?</p>
</blockquote>
</aside>
<p>I have a use case:<br>
I would like to set an attribute of the Module object that defines what layer should be used for a specific task - for instance, that resnet.fc layer should be used as a target for embedding.</p>
<p>If I specify <code>resnet.target_layer = resnet.fc</code>, the layer now appears twice in _modules.</p>
<p>Getting the layer by string name doesn’t work well for scenarios where the layer might be buried in several nests, such as <code>efficientnet.features[-1].conv2d</code> (made that up, might not exist).</p>
<p>So how should I make a “pointer” to the layer without duplicating it in <code>._modules</code>?</p>
<p>(as mentioned above, a workaround is to wrap these attributes in something like a tuple/list/dictionary)</p>

username: sammlapp
cooked: <p>since I need to be able to replace/modify the layer as well as access it, I reverted to using string descriptors for the layer. The layer can be accessed with module.get_submodule(‘feat.0.conv1’) and I wrote a helper function to parse this type of naming convention for replacing the layer. Sharing here in case its useful to others:</p>
<pre><code class="lang-auto">def set_layer_from_name(module, layer_name, new_layer):
    """assign an attribute of an object using a string name

    Args:
        module: object to assign attribute to
        layer_name: string name of attribute to assign
                the attribute_name is formatted with . delimiter and can contain
                either attribute names or list indices
                e.g. "network.classifier.0.0.fc" sets network.classifier[0][0].fc
                this type of string is given by torch.nn.Module.named_modules()
        new_layer: replace layer with this torch.nn.Module instance

        see also: torch.nn.Module.named_modules(), torch.nn.Module.get_submodule()
    """
    parts = layer_name.split(".")
    for part in parts[:-1]:
        # might be an index of a list
        try:
            part = int(part)
            module = module[part]
        except ValueError:
            module = getattr(module, part)
    try:
        part = int(parts[-1])
        module[part] = new_layer
    except ValueError:
        setattr(module, parts[-1], new_layer)
</code></pre>