Skip to content

Solve checkpoint and validation bugs in training corrdiff unet #1000

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 10 commits into
base: main
Choose a base branch
from

Conversation

albertocarpentieri
Copy link

@albertocarpentieri albertocarpentieri commented Jul 3, 2025

  • solved bug in training to enable validation
  • add use_reentrant=False to enable checkpointing and compiling in song unet
  • make dimensions consistent in song unet for checkpoint_threshold computation

Closes #994 .

@albertocarpentieri
Copy link
Author

@CharlelieLrt Can you please review it?

@CharlelieLrt CharlelieLrt self-requested a review July 3, 2025 18:12
Copy link
Collaborator

@CharlelieLrt CharlelieLrt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, just left a suggestion to make the checkpointing more robust for high/low aspect ratio images.

Comment on lines -1156 to +1154
idx_x = torch.arange(self.img_shape_y)
idx_y = torch.arange(self.img_shape_x)
mesh_x, mesh_y = torch.meshgrid(idx_x, idx_y)
grid = torch.stack((mesh_x, mesh_y), dim=0) # (2, img_shape_y, img_shape_x)
idx_x = torch.arange(self.img_shape_x)
idx_y = torch.arange(self.img_shape_y)
mesh_y, mesh_x = torch.meshgrid(idx_y, idx_x)
grid = torch.stack((mesh_y, mesh_x), dim=0) # (2, img_shape_y, img_shape_x)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your changes are consistent with the original implementation, but I think the original implementation is wrong.
If you look at the shape of grid, I think it is transposed in comparison to the case self.gridtype == "linear".

The case self.gridtype == "test" is only used in CI tests, but it might be better to fix the implementation + fix any test that is failing (probably not many because we have only very few tests for non-square images)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

🐛[BUG]: CorrDiff training with SongUnet checkpoint support
2 participants