Propose fix perceptual loss sqrt nan#8414
Propose fix perceptual loss sqrt nan#8414cvbourne wants to merge 6 commits intoProject-MONAI:devfrom
Conversation
|
Thanks for the update, the changes looks fine to me. |
| def normalize_tensor(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: | ||
| norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) | ||
| return x / (norm_factor + eps) |
There was a problem hiding this comment.
| def normalize_tensor(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: | |
| norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) | |
| return x / (norm_factor + eps) | |
| def normalize_tensor(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: | |
| norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) | |
| return x / norm_factor |
Do we want to remove eps from the denominator? As proposed eps will contribute twice to the final result.
There was a problem hiding this comment.
This file should go into an appropriate subdirectory in the tests directory. We've changed the directory structure there recently so probably tests/losses.
| # Create tensor | ||
| x = torch.zeros(2, 3, 10, 10, requires_grad=True) | ||
|
|
||
| optimizer = optim.Adam([x], lr=0.01) |
There was a problem hiding this comment.
I don't think the optimizer is needed for this test?
| x = torch.zeros(2, 3, 10, 10, requires_grad=True) | ||
|
|
||
| optimizer = optim.Adam([x], lr=0.01) | ||
| x_scaled = x * scale |
There was a problem hiding this comment.
Since x is all 0, x_scaled is always going to be 0 unless you're expected float imprecision to create values here. If so, I would add a comment to mention this.
There was a problem hiding this comment.
I don't understand the point of this test with regards to the next one; instead of a zeros tensor, couldn't it be a random one which will be then multiplied by a really small number?
virginiafdez
left a comment
virginiafdez
left a comment
There was a problem hiding this comment.
The changes look good to me. I'd modify one of the tests, but the rest is fine.
| x = torch.zeros(2, 3, 10, 10, requires_grad=True) | ||
|
|
||
| optimizer = optim.Adam([x], lr=0.01) | ||
| x_scaled = x * scale |
There was a problem hiding this comment.
I don't understand the point of this test with regards to the next one; instead of a zeros tensor, couldn't it be a random one which will be then multiplied by a really small number?
virginiafdez
left a comment
virginiafdez
left a comment
There was a problem hiding this comment.
Besides my comment about the point of one of the tests, I think this PR can be merged, as long as the errors happening on the automatic tests are fixed.
|
Hi @cvbourne, could you please help resolve the DCO issue and also help take a look at the failed pipeline? Thanks. |
📝 WalkthroughWalkthroughNumerical stability improvements to the normalize_tensor function in the perceptual loss module by increasing default epsilon from 1e-10 to 1e-8 and repositioning epsilon inside the squared sum calculation before the square root operation. A new test file validates gradient stability with small-valued and zero tensor inputs. Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
tests/test_perceptual_loss_stability.py (1)
36-39:⚠️ Potential issue | 🟡 MinorZero tensor negates scale parameter.
torch.zeros * scaleis always zeros. To test small values, use random tensor:Proposed fix
- x = torch.zeros(2, 3, 10, 10, requires_grad=True) - - optimizer = optim.Adam([x], lr=0.01) - x_scaled = x * scale + x = torch.randn(2, 3, 10, 10, requires_grad=True) + x_scaled = x * scaleThis also addresses the unused
nameparameter warning (ARG002) since parameterized tests require the name argument for test identification.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_perceptual_loss_stability.py` around lines 36 - 39, The test currently creates x = torch.zeros(...) then computes x_scaled = x * scale which is always zero; replace the zero tensor with a small random tensor (e.g., torch.randn(...) * small_factor or torch.empty(...).normal_(mean=0, std=small_value)) so scaling actually affects values, and ensure the parameterized test's name parameter is used (or include it in the test id) to avoid the unused-name warning; update references to x, x_scaled, optimizer in the test_perceptual_loss_stability test accordingly.
🧹 Nitpick comments (2)
tests/test_perceptual_loss_stability.py (1)
17-17: Remove unused import.
torch.optimimport can be removed when optimizer is deleted.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_perceptual_loss_stability.py` at line 17, Remove the now-unused import "import torch.optim as optim" from tests/test_perceptual_loss_stability.py because the optimizer was removed; search for the symbol "optim" and delete the import statement so the file no longer contains an unused dependency.monai/losses/perceptual.py (1)
277-279: Missing docstring.Per coding guidelines, add a docstring describing parameters (
x,eps) and return value.Proposed docstring
def normalize_tensor(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: + """Normalize tensor across channel dimension. + + Args: + x: Input tensor with shape (B, C, ...). + eps: Small constant for numerical stability. + + Returns: + Channel-normalized tensor with same shape as input. + """ norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) - return x / (norm_factor + eps) + return x / norm_factor🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@monai/losses/perceptual.py` around lines 277 - 279, Add a proper docstring to the normalize_tensor function that documents parameters and return value: describe param x (torch.Tensor) as the input tensor with channel dimension at dim=1, param eps (float) as the small numerical epsilon used to stabilize the norm computation, and describe the return as a torch.Tensor containing the input scaled by the L2 norm per-channel (with shape preserved). Also mention behavior with shapes and that eps prevents division by zero and the dtype of the returned tensor matches the input.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@monai/losses/perceptual.py`:
- Around line 277-279: The normalize_tensor function currently applies eps twice
(inside the sqrt and again added to norm_factor), so change normalize_tensor to
only add eps inside the sqrt: compute norm_factor = torch.sqrt(torch.sum(x**2,
dim=1, keepdim=True) + eps) and return x / norm_factor (remove the extra + eps
in the denominator) to avoid double-counting the epsilon while keeping numerical
stability.
---
Duplicate comments:
In `@tests/test_perceptual_loss_stability.py`:
- Around line 36-39: The test currently creates x = torch.zeros(...) then
computes x_scaled = x * scale which is always zero; replace the zero tensor with
a small random tensor (e.g., torch.randn(...) * small_factor or
torch.empty(...).normal_(mean=0, std=small_value)) so scaling actually affects
values, and ensure the parameterized test's name parameter is used (or include
it in the test id) to avoid the unused-name warning; update references to x,
x_scaled, optimizer in the test_perceptual_loss_stability test accordingly.
---
Nitpick comments:
In `@monai/losses/perceptual.py`:
- Around line 277-279: Add a proper docstring to the normalize_tensor function
that documents parameters and return value: describe param x (torch.Tensor) as
the input tensor with channel dimension at dim=1, param eps (float) as the small
numerical epsilon used to stabilize the norm computation, and describe the
return as a torch.Tensor containing the input scaled by the L2 norm per-channel
(with shape preserved). Also mention behavior with shapes and that eps prevents
division by zero and the dtype of the returned tensor matches the input.
In `@tests/test_perceptual_loss_stability.py`:
- Line 17: Remove the now-unused import "import torch.optim as optim" from
tests/test_perceptual_loss_stability.py because the optimizer was removed;
search for the symbol "optim" and delete the import statement so the file no
longer contains an unused dependency.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Cache: Disabled due to data retention organization setting
Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting
📒 Files selected for processing (2)
monai/losses/perceptual.pytests/test_perceptual_loss_stability.py
| def normalize_tensor(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: | ||
| norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) | ||
| return x / (norm_factor + eps) |
There was a problem hiding this comment.
Double epsilon still present.
Per past review discussion, eps contributes twice: once inside sqrt() and again in the denominator. Author agreed to remove the second one.
Proposed fix
def normalize_tensor(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps)
- return x / (norm_factor + eps)
+ return x / norm_factor🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@monai/losses/perceptual.py` around lines 277 - 279, The normalize_tensor
function currently applies eps twice (inside the sqrt and again added to
norm_factor), so change normalize_tensor to only add eps inside the sqrt:
compute norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) and
return x / norm_factor (remove the extra + eps in the denominator) to avoid
double-counting the epsilon while keeping numerical stability.
4 participants
Fixes # 8412
Description
This PR fixes a numerical stability issue in the PerceptualLoss implementation where the
normalize_tensorfunction can produce NaN gradients when the input values are very small.Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.