Adding a normalized MSE loss. (#7)

* Adding runs to .gitignore, so that no models/results are pushed to Git.

* Adding the loss function.

* Adding linting and testing to the NormalizedMSELoss.

* Fixed the tests.

* Linting tests and fixing them so that they pass.

* More linting fixes.

* Fixing linting on the losses.

* Update README.md

---------

Signed-off-by: Gal Egozi <[email protected]>

Author: galegozi

.github/workflows/yoke_install_test_lint.yml

@@ -16,64 +16,84 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+
     - name: Set up Python 3.9
       uses: actions/setup-python@v3
       with:
         python-version: '3.9'
+
     - name: Install flit
       run: pip install flit
+
     - name: Build yoke
       run: flit install --deps=all
+
     - name: Test with pytest
       run: |
         pytest -v --cov-report=lcov:./coverage/lcov.info --cov=yoke -Werror
+
     - name: Upload coverage to Coveralls
       uses: coverallsapp/github-action@v1
+
     - name: Lint Yoke
       run: |
         ruff check
         ruff check --preview
         ruff format --check --diff
       continue-on-error: true
+
     - name: Lint applications/evaluation
       run: |
         ruff check applications/evaluation
         ruff check applications/evaluation --preview
         ruff format applications/evaluation --check --diff
       continue-on-error: false
+
     - name: Lint applications/filelists
       run: |
         ruff check applications/filelists
         ruff check applications/filelists --preview
         ruff format applications/filelists --check --diff
       continue-on-error: false
+
     - name: Lint applications/normalization
       run: |
         ruff check applications/normalization
         ruff check applications/normalization --preview
         ruff format applications/normalization --check --diff
       continue-on-error: false
+
     - name: Lint applications/viewers
       run: |
         ruff check applications/viewers
         ruff check applications/viewers --preview
         ruff format applications/viewers --check --diff
       continue-on-error: false
+
     - name: Lint tests
       run: |
         ruff check tests
         ruff check tests --preview
         ruff format tests --check --diff
       continue-on-error: false
+
     - name: Lint datasets
       run: |
         ruff check src/yoke/datasets
         ruff check src/yoke/datasets --preview
         ruff format src/yoke/datasets --check --diff
       continue-on-error: false
+
     - name: Lint models
       run: |
         ruff check src/yoke/models
         ruff check src/yoke/models --preview
         ruff format src/yoke/models --check --diff
       continue-on-error: false
+
+    - name: lint losses
+      run: |
+        ruff check src/yoke/losses
+        ruff check src/yoke/losses --preview
+        ruff format src/yoke/losses --check --diff
+      continue-on-error: false

README.md

@@ -1,13 +1,12 @@
 YOKE: Yielding Optimal Knowledge Enhancement
 ============================================
 
+[![Coverage Status](https://coveralls.io/repos/github/lanl/Yoke/badge.svg?branch=main)](https://coveralls.io/github/lanl/Yoke?branch=main)
 [![pipeline status](https://github.com/lanl/Yoke/actions/workflows/yoke_install_test_lint.yml/badge.svg)](https://github.com/lanl/Yoke/actions) 
 [![Latest Release](https://img.shields.io/github/v/release/lanl/Yoke)](https://github.com/lanl/Yoke/releases)
-[![Coverage Status](https://coveralls.io/repos/github/lanl/Yoke/badge.svg?branch=main)](https://coveralls.io/github/lanl/Yoke?branch=main)
 
 ![Get YOKEd!](./YOKE_DALLE_512x512.png)
 
-
 About:
 ------
 
@@ -19,11 +18,14 @@ projects.
 The YOKE module is divided into submodules, installed in a python environment:
 
 - datasets/
+- helpers/
 - models/
 - metrics/
-- torch_training_utils.py
+- losses/
+- utils/
 - lr_schedulers.py
-- parallel_utils.py
+- parellel_utils.py
+- torch_training_utils.py
 
 Helper utilities and examples are under `applications`:
 

src/yoke/losses/NormMSE.py

@@ -0,0 +1,56 @@
+"""This loss function is a per-channel normalized version of mean squared error."""
+
+import torch
+import torch.nn as nn
+
+
+class NormalizedMSELoss(nn.Module):
+    """Per-channel normalized mean squared error loss.
+
+    This loss function normalizes the input and target tensors per channel
+    before computing the mean squared error. The normalization is done by
+    subtracting the mean and dividing by the standard deviation of the target
+    tensor, with a small epsilon added to the standard deviation to avoid
+    division by zero.
+
+    Args:
+        eps (float): A small value to avoid division by zero.
+        reduction (str): Specifies the reduction to apply to the output:
+            'none' | 'mean' | 'sum'. Default: 'none'.
+    """
+
+    def __init__(self, eps: float = 1e-8, reduction: str = "none") -> None:
+        """Initialize the NormalizedMSELoss.
+
+        Args:
+            eps (float): A small value to avoid division by zero.
+            reduction (str): Specifies the reduction to apply to the output:
+                'none' | 'mean' | 'sum'. Default: 'none'.
+        """
+        super().__init__()
+        self.eps = eps
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """Compute the normalized mean squared error loss.
+
+        Args:
+            pred (torch.Tensor): The predicted tensor.
+            target (torch.Tensor): The target tensor.
+
+        Returns:
+            torch.Tensor: The computed loss.
+        """
+        target_mean = target.mean(dim=(0, 2, 3), keepdim=True)
+        target_std = target.std(dim=(0, 2, 3), keepdim=True) + self.eps
+
+        pred_norm = (pred - target_mean) / target_std
+        target_norm = (target - target_mean) / target_std
+
+        loss = (pred_norm - target_norm) ** 2
+        if self.reduction == "mean":
+            return loss.mean()
+        elif self.reduction == "sum":
+            return loss.sum()
+        else:
+            return loss

src/yoke/losses/__init__.py

@@ -0,0 +1 @@
+"""All custom loss functions in YOKE live here."""

tests/losses/__init__.py

@@ -0,0 +1 @@
+"""The tests for the src/yoke/losses directory."""

tests/losses/test_normMSEloss.py

@@ -0,0 +1,130 @@
+"""Tests for the NormalizedMSELoss class."""
+
+import pytest
+import torch
+from yoke.losses.NormMSE import NormalizedMSELoss
+
+
+@pytest.fixture
+def norm_mse() -> NormalizedMSELoss:
+    """Fixture for NormalizedMSELoss."""
+    return NormalizedMSELoss()
+
+
+def test_norm_mse_loss_zero(norm_mse: NormalizedMSELoss) -> None:
+    """Test the NormalizedMSELoss with zero input and target."""
+    inp = torch.tensor([[[[0.0, 0.0], [0.0, 0.0]]]])
+    target = torch.tensor([[[[0.0, 0.0], [0.0, 0.0]]]])
+    loss = norm_mse(inp, target)
+    assert torch.all(loss == 0.0)
+
+
+def test_norm_mse_loss_positive(norm_mse: NormalizedMSELoss) -> None:
+    """Test the NormalizedMSELoss with positive input and target."""
+    inp = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], dtype=torch.float32)
+    target = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]], dtype=torch.float32)
+    loss = norm_mse(inp, target)
+    assert torch.all(loss == 0.0)
+
+
+def test_norm_mse_loss_non_zero(norm_mse: NormalizedMSELoss) -> None:
+    """Test the NormalizedMSELoss with non-zero input and target."""
+    inp = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
+    target = torch.tensor([[[[4.0, 5.0], [6.0, 7.0]]]])
+    loss = norm_mse(inp, target)
+    expected_loss = torch.mean(
+        (
+            (inp - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+            - (target - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+        )
+        ** 2,
+        dim=(0, 2, 3),
+    )
+    assert torch.all(loss == expected_loss)
+
+
+def test_norm_mse_loss_negative(norm_mse: NormalizedMSELoss) -> None:
+    """Test the NormalizedMSELoss with negative input and target."""
+    inp = torch.tensor([[[[-1.0, -2.0], [-3.0, -4.0]]]])
+    target = torch.tensor([[[[-4.0, -5.0], [-6.0, -7.0]]]])
+    loss = norm_mse(inp, target)
+    expected_loss = torch.mean(
+        (
+            (inp - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+            - (target - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+        )
+        ** 2,
+        dim=(0, 2, 3),
+    )
+    assert torch.all(loss == expected_loss)
+
+
+def test_norm_mse_loss_mean_reduction() -> None:
+    """Test the mean reduction of the NormalizedMSELoss."""
+    norm_mse = NormalizedMSELoss(reduction="mean")
+    inp = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
+    target = torch.tensor([[[[4.0, 5.0], [6.0, 7.0]]]])
+    loss = norm_mse(inp, target)
+    expected_loss = torch.mean(
+        (
+            (inp - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+            - (target - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+        )
+        ** 2
+    ).mean()
+    assert torch.all(loss == expected_loss)
+
+
+def test_norm_mse_loss_sum_reduction() -> None:
+    """Test the sum reduction of the NormalizedMSELoss."""
+    norm_mse = NormalizedMSELoss(reduction="sum")
+    inp = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
+    target = torch.tensor([[[[4.0, 5.0], [6.0, 7.0]]]])
+    loss = norm_mse(inp, target)
+    expected_loss = torch.sum(
+        (
+            (inp - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+            - (target - target.mean(dim=(0, 2, 3), keepdim=True))
+            / (target.std(dim=(0, 2, 3), keepdim=True) + norm_mse.eps)
+        )
+        ** 2
+    ).sum()
+    assert loss == expected_loss
+
+
+def test_norm_mse_loss_different_shapes(norm_mse: NormalizedMSELoss) -> None:
+    """Test the NormalizedMSELoss with different input and target shapes."""
+    inp = torch.tensor([[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]])
+    target = torch.tensor([[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]])
+    loss = norm_mse(inp, target)
+    assert torch.all(loss == 0.0)
+
+
+def test_norm_mse_loss_batch_size(norm_mse: NormalizedMSELoss) -> None:
+    """Test the NormalizedMSELoss with different batch sizes."""
+    inp = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]], [[[1.0, 2.0], [3.0, 4.0]]]])
+    target = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]], [[[1.0, 2.0], [3.0, 4.0]]]])
+    loss = norm_mse(inp, target)
+    assert torch.all(loss == 0.0)
+
+
+def test_norm_mse_loss_different_eps() -> None:
+    """Test the NormalizedMSELoss with different eps values."""
+    norm_mse = NormalizedMSELoss(eps=1e-5)
+    inp = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
+    target = torch.tensor([[[[4.0, 5.0], [6.0, 7.0]]]])
+    loss = norm_mse(inp, target)
+    expected_loss = (
+        (inp - target.mean(dim=(0, 2, 3), keepdim=True))
+        / (target.std(dim=(0, 2, 3), keepdim=True) + 1e-5)
+        - (target - target.mean(dim=(0, 2, 3), keepdim=True))
+        / (target.std(dim=(0, 2, 3), keepdim=True) + 1e-5)
+    ) ** 2
+    assert torch.all(loss == expected_loss)