Sometime we need to use apply_optimizer_in_backward for performance. However, in some cases when some parts are not differentiable, we have to manually calculate part of the gradients and manually apply the gradients.
For example, in the code snippet below, there are two losses, one can be directly calculated using nn.MSELoss, another has non-differentiable part and must be calculated manually.
- Without fusion, there are 2 gradients accumulation and 1 optimization step
- With fusion, there are 2 (gradient+optimize) fused step.
My question is: now I perform 2 optimize, it might affect the convergence, what is the best practice for such situation?
# linear layer and loss function
mse_loss = nn.MSELoss()
# input data and labels
x = torch.rand(3 ,5)
labels = torch.randint(low=0, high=2, size=(3, 1))
# logits
logits = model(x)
# loss
loss1 = mse_loss(logits, labels)
loss1.backward() # this will perform a backward&optimize fusion, and parameters will be updated?
# another loss with non-differentiable part, must be calculated manually
custom_gradient = NonDifferentiableLogic(...)
torch.autograd.backward(logits, grad_tensors=custom_gradient) # this will perform another backward&optimize fusion, and parameters will be updated again?