Skip to content

Resolves non-gradient bug of SSIM Loss#5686

Merged
wyli merged 15 commits intoProject-MONAI:devfrom
PedroFerreiradaCosta:dev
Dec 9, 2022
Merged

Resolves non-gradient bug of SSIM Loss#5686
wyli merged 15 commits intoProject-MONAI:devfrom
PedroFerreiradaCosta:dev

Conversation

@PedroFerreiradaCosta
Copy link
Copy Markdown
Contributor

Fixes #5668 .

Description

The SSIM loss wasn't feeding into autograd - it wasn't allowing gradient to be calculated. Now they can be calculated.
As per my previous PR (#5550), the loss function is now calculated using the SSIMMetric (and not the other way around). The SSIMMetric in inheriting from class IterationMetric that runs .detach() on all inputs before running ._compute_tensor().

y_ = y.detach() if isinstance(y, torch.Tensor) else None

I'm now calling the SSIMMetric()._compute_tensor() explicitly, removing the step where it detaches from the current graph.

To Reproduce

import torch
from monai.losses.ssim_loss import SSIMLoss

x = torch.ones([1,1,10,10])/2
y = torch.ones([1,1,10,10])/2
y.requires_grad_(True)
data_range = x.max().unsqueeze(0)
loss = SSIMLoss(spatial_dims=2)(x,y,data_range)
print(loss.requires_grad)

Expected behavior
loss.requires_grad will now be True.

Types of changes

  • Non-breaking change (fix or new feature that would not break existing functionality).

@wyli
Copy link
Copy Markdown
Contributor

wyli commented Dec 8, 2022

thanks, could you please add a unit test

@wyli
Copy link
Copy Markdown
Contributor

wyli commented Dec 8, 2022

could you please help test these changes in a training pipeline? @Can-Zhao @mersad95zd

Signed-off-by: PedroFerreiradaCosta <[email protected]>
@PedroFerreiradaCosta
Copy link
Copy Markdown
Contributor Author

@wyli submitted now a unit test in test_ssim_loss.py for checking if loss.requires_grad==True

Copy link
Copy Markdown
Collaborator

@Can-Zhao Can-Zhao left a comment

Choose a reason for hiding this comment

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

Thank you so much!!

@wyli
Copy link
Copy Markdown
Contributor

wyli commented Dec 8, 2022

/build

Copy link
Copy Markdown
Contributor

@wyli wyli left a comment

Choose a reason for hiding this comment

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

premerge tests work fine as well

@wyli
Copy link
Copy Markdown
Contributor

wyli commented Dec 9, 2022

merging this for now, please feel free to follow up if you have further code changes.

@wyli wyli marked this pull request as ready for review December 9, 2022 08:36
@wyli wyli merged commit af44576 into Project-MONAI:dev Dec 9, 2022
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.

SSIMLoss result does not have gradient

3 participants