Non-leaf variables as a module's parameters

1

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
# nn.module.weight.data = d.data  # 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!

2

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
3

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.

Thanks!

4

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:

5

Hi,

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