pytorch custom loss with regards to weights

user_12918011 0
Viewed 426

I have a loss based on 2 things:

  1. MSE loss
  2. A custom loss term based on the network weights.

I have this code:

net = CustomNet()
mse_loss = torch.nn.MSELoss()
def custom_loss(output, target):
    weights = net.linear_layer.weight.data
    return mse_loss(output, target) + torch.linalg.norm(weights @ weights.T -
                                                        torch.eye(weights.shape[0]))

When I try to remove the MSE loss (so my loss is only based on the weights):

def custom_loss(output, target):
    weights = net.linear_layer.weight.data
    return torch.linalg.norm(weights @ weights.T -
                             torch.eye(weights.shape[0]))

I am getting the error:

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I can notice that mse loss has grad_fn=<MseLossBackward object at 0x14908c450>

What am I doing wrong? Why can't I use only the second loss?

pytorch
1 Answers

You can't use the second term alone because it doesn't have a grad_fn function as the first term does. This means if you're having both terms it will only backpropagate on the first term (the MSE loss) and will not consider the second. Having no grad_fn means it is seen as a constant w.r.t. the input or parameter and has no effect on the gradient.

The tensors you use to compute the second term do not require a gradient. More specifically any tensor that you get using the data attribute won't require a gradient. In your case net.linear_layer.weight.data.

Instead you should access the tensor directly via:

>>> weights = net.linear_layer.weight
user_63313690
Related