I am working on a project that requires me to write a method that is not differentiable. Hence I have calculated the gradient w.r.t the input of this method. Now I have tried to use register_hook to return the calculated grad to be able to flow the gradient backwards from there. But register_hook is not being called for this tensor. After reading through a few posts here, I realized that register_hook will not be called for tensors if gradient is not being calculated. I tried to assign gradient by doing this. my_tensor.grad = calculated_grad. However, I do not think that has any effect.
Can you please help me on this? I am sharing a demo version of my code.
class MLP(nn.Module):
def __init__(self, input_size, hidden_size, num_classes=2):
super(MLP, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.fc1 = nn.Linear(self.input_size, self.hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(self.hidden_size,num_classes)
def forward(self, x):
hidden = self.fc1(x)
relu = self.relu(hidden)
output = self.fc2(relu)
return output
net = MLP(input_size, hidden_size, num_classes)
out = net(data)
#used out on a non-differentiable method
final_output = non-diff(out)
calculated_grad = grad_calc(out) # calculated gradient of out with respect to non-diff method.
loss = criterion(final_output, target)
loss.backward()
out.grad = calculated_grad
opt.step()
I would really appreciate your help on this. I have tried so many things to fix this. But nothing worked.
