MouseLand
[BUG] semantic error in learning objective loss function
I think there is a sematic error in the learning objective loss function. in the train_seg function, at the part where you would compute mse between predicted flows and Gt flows:
`def _loss_fn_seg(lbl, y, device): """ Calculates the loss function between true labels lbl and prediction y.
Args:
lbl (numpy.ndarray): True labels (cellprob, flowsY, flowsX).
y (torch.Tensor): Predicted values (flowsY, flowsX, cellprob).
device (torch.device): Device on which the tensors are located.
Returns:
torch.Tensor: Loss value.
"""
criterion = nn.MSELoss(reduction="mean")
criterion2 = nn.BCEWithLogitsLoss(reduction="mean")
veci = 5. * lbl[:, -2:]
loss = criterion(y[:, -3:-1], veci)
loss /= 2.
loss2 = criterion2(y[:, -1], (lbl[:, -3] > 0.5).float())
loss = loss + loss2
return loss
`
notice that in veci = 5. * lbl[:, -2:] loss = criterion(y[:, -3:-1], veci), you compute mse bewteen GT and prediction. GT flows are multiplied by a factor of 5, while the prediction is not. Mathematically, you are multiplying only one element of the substraction and not the square of the substraction which I think would be correct. The paper claims to scale the mse to a factor of 5 to match the bce loss but I dont think this is the correct way.
Btw Mediar did the same. So either they just blindly copied their semantic or I am missing something.
@le1nax I'm not sure where the paper says that the mse is scaled, but the relevant equation from the cellpose 1 paper is
which is what is implemented in the code.
@le1nax I'm not sure where the paper says that the mse is scaled, but the relevant equation from the cellpose 1 paper is
which is what is implemented in the code.
Hello, Thank you very much for your answer and your time.
Yes that is indeed the equation that is implemented in the code, I was never questioning this.
My concern is about the equation itself:
First, lets clarify what is actually the intention behind the scaling of the ground truth flow:
To me, it seems logical to balance the loss contributions from the flows and the cell probabilities. This is, what I also assume the paper means:
"To match the relative contributions of the L2 loss and cross-entropy loss, we multiplied the gradients by a factor of five."
This seems to me like balancing the numerical loss contributions, which to me, is totally logical: For instance, when logging the flow L2 loss and the probability BCE loss, one could see outputs like: Flow loss: 0.94, Probability loss: 0.25. In this case it seems logical to multiply the BCE loss contribution by 4, to match it with the Flow loss.
This is of course just an example. I assume the factor 5 was chosen because the flow loss was about 1/5 of the BCE loss, over multiple consecutive iterations.
Now, when scaling the loss, to me it only seems logical to first calculate each loss component and then scale the component by a scalar:
Loss_flow = mean(Flow_pred^2 - Flow_gt^2) Loss_prob = BCE(sigmoid(Prob_pred), Prob_gt)
Loss_total = Loss_flow * 5 + Loss_prob
Instead, what cellpose does is:
Loss_flow = mean(Flow_pred - 5* Flow_gt)^2 Loss_prob = BCE(sigmoid(Prob_pred), Prob_gt)
Loss_total = Loss_flow + Loss_prob
Note that the actual loss components are not scaled, but the Flows inside the substraction, which to me seems incorrect.
As my issue was closed due to duplicated and it seems like it's not been answered clearly here. What is the fundamental choice of forcing the model to predict flow 5 times greater than the labels? As @le1nax mentionned, scaling the loss to match another loss's scale in the multi-objective would be something expected but this is not what's been done in cellpose.
"flow" does not have an intrinsic unit of measurement. Multiplying by 5 was just a choice of scaling, which we have kept in order to be backwards compatible. If you prefer, in your own code you can predict unscaled flows, or you can scale them (and the loss functions) however you like. The Cellpose flows will always be scaled by 5 for backwards compatibility.
