Lightning-AI / Lightning-AI/pytorch-lightning
loss spikes in validation step when the model has multiple losses applied
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
I have a model that has multiple losses applied. see the code following:
```
# [B, H, W, 128]
x_bin_feature=F.softmax(x_bin_pred, dim=3)
# [B, H, W, 128]
y_bin_feature=F.softmax(y_bin_pred, dim=3)
# [B, H, W, 128]
z_bin_feature=F.softmax(z_bin_pred, dim=3)
# [B, H, W, 384]
original_size_feature_map=torch.cat([x_bin_feature, y_bin_feature, z_bin_feature], dim=3)
# [B, 384, H, W]
original_size_feature_map = torch.permute(
original_size_feature_map, (0, 3, 1, 2)
)
pnp_feature = self._pnp_head(original_size_feature_map)
pnp_feature = pnp_feature.reshape(pnp_feature.shape[0], -1)
pnp_feature = self._pnp_neck(pnp_feature)
outputs["trans"] = self._pnp_trans_head(pnp_feature)
outputs["rot"] = self._pnp_rot_head(pnp_feature)
coordinate_x_loss = self._coordinate_x_loss(
x_bin_pred.view(-1, 128),
x_bin_gt.long().view(-1),
)
coordinate_y_loss = self._coordinate_y_loss(
y_bin_pred.view(-1, 128),
y_bin_gt.long().view(-1),
)
coordinate_z_loss = self._coordinate_z_loss(
z_bin_pred.view(-1, 128),
z_bin_gt.long().view(-1),
)
coordinate_loss=coordinate_x_loss+coordinate_y_loss+coordinate_z_loss
trans_label, rot_label = stack_pose_labels(pose_labels)
pm_loss = self._pm_loss(
rotation_matrix_preds,
rotation_matrix_gt,
all_points,
outputs["trans"],
trans_label,
)
```
I used the same validation set as the training set. And I train the model as usual, however I saw that the losses consistently dropping in training steps, but everytime in validation step, the loss goes extremely high. And after the validation steps, the training loss drops with no influence of validation steps:

When I remove one of the loss in my code, the model can be trained normally.
Is there sth wrong with how I use the multiple losses? Why the training loss is so different with the validation loss?
### What version are you seeing the problem on?
v1.x
### How to reproduce the bug
_No response_
### Error messages and logs
```
# Error messages and logs here please
```
### Environment
Current environment
```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```
### More info
_No response_
Contributor guide
Research direction
No source files, tests, reproducible steps, or complete environment are provided. Start by reproducing the reported training and validation loss discrepancy from the supplied multi-loss code, then establish a minimal case and define done as identifying whether the behavior is in the framework or the model setup.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100