Lightning-AI / Lightning-AI/pytorch-lightning

loss spikes in validation step when the model has multiple losses applied

Open
#20,173 1 comment 0 reactions 0 assignees View on GitHub
question
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:
![image](https://github.com/user-attachments/assets/87b15f29-6db7-4898-a9ea-8aa60e445dc4)

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.