Non-leaf variables as a module's parameters

I wonder how to set up an NN module whose parameters are non-leaf variables.
What I would like to do is something like below:

class NNTest(torch.nn.Module):
    def __init__(self):
        super(NNTest, self).__init__()
        self.module = torch.nn.Linear(2, 3, bias=False)
    def forward(self, x):
        return self.module(x)

nn = NNTest()
p = torch.nn.Parameter(torch.randn(2, 3))
d = p**2
nn.module.weight = d  # raises a TypeError here
# =  # this is not what I want

The problem of the last line is that it does not propagate the gradient to p and I would like the gradient to propagate beyond the Module’s weight.
Any help will be appreciated! Thanks!

You can do this by first deleting the original one and then assigning your non-leaf Tensor:

del nn.module.weight
nn.module.weight = d
1 Like

Is there a way to do this del call on the parameters dynamically, without having to write the nested object name nn.module.weight?

Example: Say the module has N total parameter sets, and I want to delete all of them, and cannot know what they will be named pre run, because I am writing a general wrapper for any torch.nn.Module.

I do not mind reading the names at runtime and dealing with those, but I just cannot write less general code with their names explicitly typed out.

For recursive set and get, this stackoverflow post was helpful, and also I could use the .children() or .modules() functions. But I am stuck on deletion.


Update: I managed to use python’s exec function to do exec(f'del {model_name}.{module_name}') where module_name is one of the names from my model’s named_parameters(), but would appreciate any less hacky way to do it. :slight_smile:


Another way to do it is described in this other post: How does one have the parameters of a model NOT BE LEAFS?

1 Like