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

1

Hi,

I have a parent nn.Module parent that saves multiple custom defined child nn.Modules as attributes of which I want one to not appear in parent.children().

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?

Thanks!

2

If you are using a custom class (or just a Python object) for the children, I would guess that they won’t be registered in the parent nn.Module.

3

Hi, thanks for the answer!

As the modules arise systematically (and the relevant attribute is expected to be of type Module elsewhere), it would be nice to not have to wrap the child Module.

I had hoped there was some possibility to enter __setattr__ with the child Module considered to be an object of type object, or would this still make the parent Module register the child in children()?

Thanks again!

4

I don’t think there is a clean approach to still use an nn.Module, but disallow registering it inside the parent. You could check this code and try to see, if changing the self._modules attribute might work, but note that this is an internal attribute and I don’t know what unwanted side effects you might trigger.
Maybe you could explain your use case a bit more so that we could think about another approach?

5

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:

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())

Fortunately, my_module.parameters() doesn’t seem to contain duplicate parameters. But it seems cleaner (at least to me) to not have both versions of the same m registered as children.

6

One quick and dirty hack I did was to wrap the submodule in tuple or list. For example:

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())

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 .to() (for example).

1 Like
7

I have a use case:
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.

If I specify resnet.target_layer = resnet.fc, the layer now appears twice in _modules.

Getting the layer by string name doesn’t work well for scenarios where the layer might be buried in several nests, such as efficientnet.features[-1].conv2d (made that up, might not exist).

So how should I make a “pointer” to the layer without duplicating it in ._modules?

(as mentioned above, a workaround is to wrap these attributes in something like a tuple/list/dictionary)

8

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:

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)