What I am going to do is modifying weight in Conv2d after loss.backward() and before optimizer.step(). One solution is to modify weights in corresponding layers after loss.backward(). I just want to make a custom layer to do this keeping train() function clean. Here is the code snippet:
class CustomConv(nn.Conv2d): def __init__(self, *kargs, **kwargs): super(CustomConv, self).__init__(*kargs, **kwargs) def forward(self, input): input = WeightModifier(self.weight)(input) out = nn.functional.conv2d( input, self.weight, None, self.stride, self.padding, self.dilation, self.groups) return out class WeightModifier(torch.autograd.Function): def __init__(self, weight): self.weight = weight def forward(self, input): return input.clone() def backward(self, grad): # Some op to change self.weight. # In other words, change weights in the following conv2d # after loss.backward() and before optimizer.step() return grad.clone() hidden = CustomConv()(input_tensor) # backward not called hidden = CustomConv()(hidden) loss = cal_loss(hidden, gt) optimizer.zero_grad() loss.backward() optimizer.step()
The problem is the backward() of WeightModifier in the first CustomConv is not called (that in second CustomConv is called). I guess the reason is that Pytorch finds that input_tensor does not require gradient and layer WeightModifier does not have any parameters. Is the any method to force or “cheat” Pytorch to execute the backward()?