[FIXED] Single weight definition and freezing in Conv2d

Issue

I am dealing with a Conv2d(3,3,kernel_size=5, stride=1) and I'd need to set some specific weights to zero and render them non updatable. For example, if I do something like

model = nn.Conv2d(3,3,kernel_size=5, stride=1)

model.weight.requires_grad = False

everything works but it affects the whole layer. I'd want to do something like this instead:

model = nn.Conv2d(3,3,kernel_size=5, stride=1)

model.weight[0,2].requires_grad = False  # this line does not work
model.weight[0,2] = 0                    # this line does not work either

it just does not seem to support assignment and requires_grad manipulation for layer parameter subgroups. Has anyone already tackled this issue?

Solution

You can zero out this filter channel by either accessing the data attribute

>>> model.weight.data[0, 2] = 0

or using the torch.no_grad context manager:

>>> with torch.no_grad():
...     model.weight[0, 2] = 0

As you noticed you can't have requires_grads specifically set for submodules. As such all parameter elements of a given module share the same flag, they either get updated or they don't.

One alternative way would be to kill the gradient of that channel manually just after the backward pass has been called, just before the optimizer step:

>>> model(torch.rand(2, 3, 100, 100)).mean().backward()
>>> model.weight.grad[0, 2] = 0
>>> optim.step()

This way your 3rd channel on filter n°1 won't get updated by the backward pass and remain at 0.

Answered By - Ivan