I am writing a NN in pytorch and I want to add the derivative of the output with respect to one of the inner layers in the loss. Here is a simple example of what I mean:
class NN(torch.nn.Module): def __init__(self): super(Model, self).__init__() self.l1 = torch.nn.Linear(2, 3) self.l2 = torch.nn.Linear(3, 1) self.l3 = torch.nn.Linear(1, 3) self.l4 = torch.nn.Linear(3, 2) def forward(self, x): z = torch.sigmoid(self.l1(x)) z = torch.sigmoid(self.l2(z)) y_pred = torch.sigmoid(self.l3(z)) y_pred = self.l4(y_pred) return y_pred
In this case (for just one example) I want the loss to be defined as:
loss = (y-y_pred)**2+norm(d(y_pred)/dz) (the last part is not correct from a python point of view). So what I want is to add the derivative of of the output with respect to a given layer to the loss function. I tried some things with hook, but that is not really what I need here. Can someone help me? Thank you!