I want to do the following operations:
import torch.nn as nn
import torch.nn.functional as F
During forward:
I defined a new conv layer:
class new_conv(nn.conv2d):
>def forward(self, input):
>>self.weight = self.weight.round()
output = F.conv2d(input, self.weight, self.bias…)
During backward:
I want to override the gradient of self.weight.round() with identity mapping, that is to replace the gradient(self.weight.round()) with gradient(self.weight).
One possible way I know is to use register_backward_hook(), however I don’t know how to apply it in my case. In Tensorflow, simply using the function G.gradient_override_map({“Round”: “Identity”}) works well.