Skip to content

fix(losses): BarlowTwinsLoss preserves float16/bfloat16 input dtype #9103

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.

Sign up for GitHub

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
tritsystem wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

fix(losses): BarlowTwinsLoss preserves float16/bfloat16 input dtype #9103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion monai/losses/barlow_twins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
c = torch.mm(input_norm.t(), target_norm) / batch_size # input_norm.t() is FxB, target_norm is BxF so c is FxF

# loss
c_diff = (c - torch.eye(c.size(0), device=c.device)).pow_(2) # FxF
c_diff = (c - torch.eye(c.size(0), dtype=c.dtype, device=c.device)).pow_(2) # FxF
c_diff[~torch.eye(c.size(0), device=c.device).bool()] *= lambd_tensor

return c_diff.sum()
11 changes: 11 additions & 0 deletions tests/losses/test_barlow_twins_loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,17 @@ def check_warning_raised(self):
with self.assertWarns(Warning):
BarlowTwinsLoss(lambd=5e-3, batch_size=1)

@parameterized.expand([(torch.float64,), (torch.float32,), (torch.bfloat16,), (torch.float16,)])
def test_preserves_input_dtype(self, dtype):
# The cross-correlation matrix `c` follows the input dtype, but `c - torch.eye(...)`
# used to silently upcast to float32 whenever dtype < float32, because torch.eye()
# was never given an explicit dtype and defaults to the global default (float32).
loss = BarlowTwinsLoss(lambd=5e-3)
i = torch.randn(4, 8, dtype=dtype)
j = torch.randn(4, 8, dtype=dtype)
output = loss(i, j)
self.assertEqual(output.dtype, dtype)


if __name__ == "__main__":
unittest.main()