deepinv 0.1.0.dev0__tar.gz → 0.2.0__tar.gz

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/PKG-INFO

@@ -1,8 +1,7 @@
 Metadata-Version: 2.1
 Name: deepinv
-Version: 0.1.0.dev0
+Version: 0.2.0
 Summary: Pytorch library for solving inverse problems with deep learning
-Author: Matthieu Terris, Samuel Hurault, Dongdong Chen
 Author-email: Julian Tachella <tachellajulian@gmail.com>
 License: BSD 3-Clause
 Project-URL: Homepage, https://deepinv.github.io/
@@ -15,11 +14,6 @@ Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: BSD License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Utilities
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development :: Libraries
@@ -29,11 +23,11 @@ License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: matplotlib
 Requires-Dist: hdf5storage
+Requires-Dist: tqdm
 Requires-Dist: torch
 Requires-Dist: torchvision
 Requires-Dist: einops
 Requires-Dist: wandb
-Requires-Dist: fastmri
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
@@ -43,9 +37,15 @@ Requires-Dist: sphinx; extra == "doc"
 Requires-Dist: sphinx_gallery; extra == "doc"
 Requires-Dist: sphinx_rtd_theme; extra == "doc"
 Requires-Dist: sphinxemoji; extra == "doc"
+Requires-Dist: sphinx_copybutton; extra == "doc"
+Requires-Dist: sphinx_autoapi; extra == "doc"
 Provides-Extra: denoisers
 Requires-Dist: bm3d; extra == "denoisers"
 Requires-Dist: timm; extra == "denoisers"
+Requires-Dist: PyWavelets; extra == "denoisers"
+Requires-Dist: ptwt; extra == "denoisers"
+Requires-Dist: FrEIA; extra == "denoisers"
+Requires-Dist: pyiqa; extra == "denoisers"
 
 .. image:: https://github.com/deepinv/deepinv/raw/main/docs/source/figures/deepinv_logolarge.png
    :width: 500px
@@ -58,18 +58,18 @@ Requires-Dist: timm; extra == "denoisers"
 
 Introduction
 ------------
-Deep Inverse is an open-source pytorch library for solving imaging inverse problems using deep learning. The goal of ``deepinv`` is to accelerate the development of deep learning based methods for imaging inverse problems, by combining popular learning-based reconstruction approaches in a common and simplified framework, standarizing forward imaging models and simplifying the creation of imaging datasets.
+Deep Inverse is an open-source pytorch library for solving imaging inverse problems using deep learning. The goal of ``deepinv`` is to accelerate the development of deep learning based methods for imaging inverse problems, by combining popular learning-based reconstruction approaches in a common and simplified framework, standardizing forward imaging models and simplifying the creation of imaging datasets.
 
-With ``deepinv`` you can:
+``deepinv`` features
 
 
-* Large collection of `predefined imaging operators <https://deepinv.github.io/deepinv/deepinv.physics.html>`_ (MRI, CT, deblurring, inpainting, etc.)
-* `Training losses <https://deepinv.github.io/deepinv/deepinv.loss.html>`_ for inverse problems (self-supervised learning, regularization, etc.).
-* Many `pretrained deep denoisers <https://deepinv.github.io/deepinv/deepinv.models.html>`_ which can be used for `plug-and-play restoration <https://deepinv.github.io/deepinv/deepinv.pnp.html>`_.
-* Framework for `building datasets <https://deepinv.github.io/deepinv/deepinv.datasets.html>`_ for inverse problems.
-* Easy-to-build `unfolded architectures <https://deepinv.github.io/deepinv/deepinv.unfolded.html>`_ (ADMM, forward-backward, deep equilibrium, etc.).
-* `Sampling algorithms <https://deepinv.github.io/deepinv/deepinv.sampling.html>`_ for uncertainty quantification (Langevin, diffusion, etc.).
-* A large number of well-explained `examples <https://deepinv.github.io/deepinv/auto_examples/index.html>`_, from basics to state-of-the-art methods.
+* A large collection of `predefined imaging operators <https://deepinv.github.io/deepinv/deepinv.physics.html>`_ (MRI, CT, deblurring, inpainting, etc.)
+* `Training losses <https://deepinv.github.io/deepinv/deepinv.loss.html>`_ for inverse problems (self-supervised learning, regularization, etc.)
+* Many `pretrained deep denoisers <https://deepinv.github.io/deepinv/deepinv.models.html>`_ which can be used for `plug-and-play restoration <https://deepinv.github.io/deepinv/deepinv.pnp.html>`_
+* A framework for `building datasets <https://deepinv.github.io/deepinv/deepinv.datasets.html>`_ for inverse problems
+* Easy-to-build `unfolded architectures <https://deepinv.github.io/deepinv/deepinv.unfolded.html>`_ (ADMM, forward-backward, deep equilibrium, etc.)
+* `Sampling algorithms <https://deepinv.github.io/deepinv/deepinv.sampling.html>`_ for uncertainty quantification (Langevin, diffusion, etc.)
+* A large number of well-explained `examples <https://deepinv.github.io/deepinv/auto_examples/index.html>`_, from basics to state-of-the-art methods
 
 .. image:: https://github.com/deepinv/deepinv/raw/main/docs/source/figures/deepinv_schematic.png
    :width: 1000px
@@ -103,22 +103,22 @@ Try out the following plug-and-play image inpainting example:
 
 .. code-block:: python
 
-    import deepinv as dinv
-    from deepinv.utils import load_url_image
-
-    url = ("https://huggingface.co/datasets/deepinv/images/resolve/main/cameraman.png?download=true")
-    x = load_url_image(url=url, img_size=512, grayscale=True, device='cpu')
-
-    physics = dinv.physics.Inpainting((1, 512, 512), mask = 0.5, \
-                                       noise_model=dinv.physics.GaussianNoise(sigma=0.01))
-
-    data_fidelity = dinv.optim.data_fidelity.L2()
-    prior = dinv.optim.prior.PnP(denoiser=dinv.models.MedianFilter())
-    model = dinv.optim.optim_builder(iteration="HQS", prior=prior, data_fidelity=data_fidelity, \
-                                     params_algo={"stepsize": 1.0, "g_param": 0.1, "lambda": 2.})
-    y = physics(x)
-    x_hat = model(y, physics)
-    dinv.utils.plot([x, y, x_hat], ["signal", "measurement", "estimate"], rescale_mode='clip')
+    >>> import deepinv as dinv
+    >>> from deepinv.utils import load_url_image
+    >>>
+    >>> url = ("https://huggingface.co/datasets/deepinv/images/resolve/main/cameraman.png?download=true")
+    >>> x = load_url_image(url=url, img_size=512, grayscale=True, device='cpu')
+    >>>
+    >>> physics = dinv.physics.Inpainting((1, 512, 512), mask = 0.5, \
+    >>>                                    noise_model=dinv.physics.GaussianNoise(sigma=0.01))
+    >>>
+    >>> data_fidelity = dinv.optim.data_fidelity.L2()
+    >>> prior = dinv.optim.prior.PnP(denoiser=dinv.models.MedianFilter())
+    >>> model = dinv.optim.optim_builder(iteration="HQS", prior=prior, data_fidelity=data_fidelity, \
+    >>>                                 params_algo={"stepsize": 1.0, "g_param": 0.1})
+    >>> y = physics(x)
+    >>> x_hat = model(y, physics)
+    >>> dinv.utils.plot([x, y, x_hat], ["signal", "measurement", "estimate"], rescale_mode='clip')
 
 
 Also try out `one of the examples <https://deepinv.github.io/deepinv/auto_examples/index.html>`_ to get started.

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/README.rst

@@ -9,18 +9,18 @@
 
 Introduction
 ------------
-Deep Inverse is an open-source pytorch library for solving imaging inverse problems using deep learning. The goal of ``deepinv`` is to accelerate the development of deep learning based methods for imaging inverse problems, by combining popular learning-based reconstruction approaches in a common and simplified framework, standarizing forward imaging models and simplifying the creation of imaging datasets.
+Deep Inverse is an open-source pytorch library for solving imaging inverse problems using deep learning. The goal of ``deepinv`` is to accelerate the development of deep learning based methods for imaging inverse problems, by combining popular learning-based reconstruction approaches in a common and simplified framework, standardizing forward imaging models and simplifying the creation of imaging datasets.
 
-With ``deepinv`` you can:
+``deepinv`` features
 
 
-* Large collection of `predefined imaging operators <https://deepinv.github.io/deepinv/deepinv.physics.html>`_ (MRI, CT, deblurring, inpainting, etc.)
-* `Training losses <https://deepinv.github.io/deepinv/deepinv.loss.html>`_ for inverse problems (self-supervised learning, regularization, etc.).
-* Many `pretrained deep denoisers <https://deepinv.github.io/deepinv/deepinv.models.html>`_ which can be used for `plug-and-play restoration <https://deepinv.github.io/deepinv/deepinv.pnp.html>`_.
-* Framework for `building datasets <https://deepinv.github.io/deepinv/deepinv.datasets.html>`_ for inverse problems.
-* Easy-to-build `unfolded architectures <https://deepinv.github.io/deepinv/deepinv.unfolded.html>`_ (ADMM, forward-backward, deep equilibrium, etc.).
-* `Sampling algorithms <https://deepinv.github.io/deepinv/deepinv.sampling.html>`_ for uncertainty quantification (Langevin, diffusion, etc.).
-* A large number of well-explained `examples <https://deepinv.github.io/deepinv/auto_examples/index.html>`_, from basics to state-of-the-art methods.
+* A large collection of `predefined imaging operators <https://deepinv.github.io/deepinv/deepinv.physics.html>`_ (MRI, CT, deblurring, inpainting, etc.)
+* `Training losses <https://deepinv.github.io/deepinv/deepinv.loss.html>`_ for inverse problems (self-supervised learning, regularization, etc.)
+* Many `pretrained deep denoisers <https://deepinv.github.io/deepinv/deepinv.models.html>`_ which can be used for `plug-and-play restoration <https://deepinv.github.io/deepinv/deepinv.pnp.html>`_
+* A framework for `building datasets <https://deepinv.github.io/deepinv/deepinv.datasets.html>`_ for inverse problems
+* Easy-to-build `unfolded architectures <https://deepinv.github.io/deepinv/deepinv.unfolded.html>`_ (ADMM, forward-backward, deep equilibrium, etc.)
+* `Sampling algorithms <https://deepinv.github.io/deepinv/deepinv.sampling.html>`_ for uncertainty quantification (Langevin, diffusion, etc.)
+* A large number of well-explained `examples <https://deepinv.github.io/deepinv/auto_examples/index.html>`_, from basics to state-of-the-art methods
 
 .. image:: https://github.com/deepinv/deepinv/raw/main/docs/source/figures/deepinv_schematic.png
    :width: 1000px
@@ -54,22 +54,22 @@ Try out the following plug-and-play image inpainting example:
 
 .. code-block:: python
 
-    import deepinv as dinv
-    from deepinv.utils import load_url_image
-
-    url = ("https://huggingface.co/datasets/deepinv/images/resolve/main/cameraman.png?download=true")
-    x = load_url_image(url=url, img_size=512, grayscale=True, device='cpu')
-
-    physics = dinv.physics.Inpainting((1, 512, 512), mask = 0.5, \
-                                       noise_model=dinv.physics.GaussianNoise(sigma=0.01))
-
-    data_fidelity = dinv.optim.data_fidelity.L2()
-    prior = dinv.optim.prior.PnP(denoiser=dinv.models.MedianFilter())
-    model = dinv.optim.optim_builder(iteration="HQS", prior=prior, data_fidelity=data_fidelity, \
-                                     params_algo={"stepsize": 1.0, "g_param": 0.1, "lambda": 2.})
-    y = physics(x)
-    x_hat = model(y, physics)
-    dinv.utils.plot([x, y, x_hat], ["signal", "measurement", "estimate"], rescale_mode='clip')
+    >>> import deepinv as dinv
+    >>> from deepinv.utils import load_url_image
+    >>>
+    >>> url = ("https://huggingface.co/datasets/deepinv/images/resolve/main/cameraman.png?download=true")
+    >>> x = load_url_image(url=url, img_size=512, grayscale=True, device='cpu')
+    >>>
+    >>> physics = dinv.physics.Inpainting((1, 512, 512), mask = 0.5, \
+    >>>                                    noise_model=dinv.physics.GaussianNoise(sigma=0.01))
+    >>>
+    >>> data_fidelity = dinv.optim.data_fidelity.L2()
+    >>> prior = dinv.optim.prior.PnP(denoiser=dinv.models.MedianFilter())
+    >>> model = dinv.optim.optim_builder(iteration="HQS", prior=prior, data_fidelity=data_fidelity, \
+    >>>                                 params_algo={"stepsize": 1.0, "g_param": 0.1})
+    >>> y = physics(x)
+    >>> x_hat = model(y, physics)
+    >>> dinv.utils.plot([x, y, x_hat], ["signal", "measurement", "estimate"], rescale_mode='clip')
 
 
 Also try out `one of the examples <https://deepinv.github.io/deepinv/auto_examples/index.html>`_ to get started.

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/__about__.py

@@ -4,14 +4,12 @@ __all__ = [
     "__url__",
     "__version__",
     "__author__",
-    "__email__",
     "__license__",
 ]
 
 __title__ = "deepinv"
 __summary__ = "Deep Learning for Inverse Problems Library for PyTorch"
-__url__ = "https://github.com/edongdongchen/deepinv"
-__version__ = "0.0.1"
-__author__ = "Dongdong Chen, Julian Tachella"
-__email__ = "echendongdong@gmail.com"
+__version__ = "0.2.0"
+__author__ = "Julian Tachella, Samuel Hurault, Matthieu Terris, Dongdong Chen"
 __license__ = "BSD 3-Clause Clear"
+__url__ = "https://deepinv.github.io/"

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/__init__.py

@@ -7,7 +7,6 @@ __all__ = [
     "__url__",
     "__version__",
     "__author__",
-    "__email__",
     "__license__",
 ]
 
@@ -55,17 +54,9 @@ from deepinv import unfolded
 
 __all__ += ["unfolded"]
 
-from deepinv.training_utils import train, test
+from deepinv.training import train, test, Trainer
+
+__all__ += ["training"]
 
 # GLOBAL PROPERTY
 dtype = torch.float
-
-# if torch.cuda.is_available():
-#    try:
-#        free_gpu_id = get_freer_gpu()
-#        device = torch.device(f"cuda:{free_gpu_id}")
-#    except:
-#        device = torch.device("cuda")
-#        print("unable to get GPU info")
-# else:
-#    device = "cpu"

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/datasets/__init__.py

@@ -1 +1,2 @@
 from .datagenerator import generate_dataset, HDF5Dataset
+from .patch_dataset import PatchDataset

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/datasets/datagenerator.py

@@ -56,6 +56,8 @@ def generate_dataset(
     batch_size=4,
     num_workers=0,
     supervised=True,
+    verbose=True,
+    show_progress_bar=False,
 ):
     r"""
     Generates dataset of signal/measurement pairs from base dataset.
@@ -88,6 +90,9 @@ def generate_dataset(
     :param bool supervised: Generates supervised pairs (x,y) of measurements and signals.
         If set to ``False``, it will generate a training dataset with measurements only (y)
         and a test dataset with pairs (x,y)
+    :param bool verbose: Output progress information in the console.
+    :param bool show_progress_bar: Show progress bar during the generation
+        of the dataset (if verbose is set to True).
 
     """
     if os.path.exists(os.path.join(save_dir, dataset_filename)):
@@ -152,17 +157,14 @@ def generate_dataset(
             if supervised:
                 hf.create_dataset("x_train", (n_train_g,) + x.shape[1:], dtype="float")
 
-            if G > 1:
-                print(
-                    f"Computing train measurement vectors from base dataset of operator {g + 1} out of {G}..."
-                )
-            else:
-                print("Computing train measurement vectors from base dataset...")
-
             index = 0
 
             epochs = int(n_train_g / len(train_dataset)) + 1
-            for e in tqdm(range(epochs)):
+            for e in (progress_bar := tqdm(range(epochs), ncols=150, disable=(not verbose or not show_progress_bar))):
+
+                desc = f"Generating dataset operator {g + 1}" if G > 1 else "Generating train dataset"
+                progress_bar.set_description(desc)
+
                 train_dataloader = DataLoader(
                     Subset(
                         train_dataset,
@@ -173,7 +175,11 @@ def generate_dataset(
                     pin_memory=False if device == "cpu" else True,
                 )
 
-                for i, x in enumerate(train_dataloader):
+                batches = len(train_dataloader) - int(train_dataloader.drop_last)
+                iterator = iter(train_dataloader)
+                for _ in range(batches):
+
+                    x = next(iterator)
                     x = x[0] if isinstance(x, list) or isinstance(x, tuple) else x
                     x = x.to(device)
 
@@ -186,9 +192,9 @@ def generate_dataset(
                     if bsize + index > n_train_g:
                         bsize = n_train_g - index
 
-                    hf["y_train"][index : index + bsize] = y[:bsize, :].to("cpu").numpy()
+                    hf["y_train"][index: index + bsize] = y[:bsize, :].to("cpu").numpy()
                     if supervised:
-                        hf["x_train"][index : index + bsize] = (
+                        hf["x_train"][index: index + bsize] = (
                             x[:bsize, :, :, :].to("cpu").numpy()
                         )
                     index = index + bsize
@@ -204,14 +210,11 @@ def generate_dataset(
                 pin_memory=True,
             )
 
-            if G > 1:
-                print(
-                    f"Computing test measurement vectors from base dataset of operator {g + 1} out of {G}..."
-                )
-            else:
-                print("Computing test measurement vectors from base dataset...")
+            batches = len(test_dataloader) - int(test_dataloader.drop_last)
+            iterator = iter(test_dataloader)
+            for i in range(batches):
 
-            for i, x in enumerate(tqdm(test_dataloader)):
+                x = next(iterator)
                 x = x[0] if isinstance(x, list) or isinstance(x, tuple) else x
                 x = x.to(device)
 
@@ -228,8 +231,8 @@ def generate_dataset(
 
                 # Add new data to it
                 bsize = x.size()[0]
-                hf["x_test"][index : index + bsize] = x.to("cpu").numpy()
-                hf["y_test"][index : index + bsize] = y.to("cpu").numpy()
+                hf["x_test"][index: index + bsize] = x.to("cpu").numpy()
+                hf["y_test"][index: index + bsize] = y.to("cpu").numpy()
                 index = index + bsize
         hf.close()
 

deepinv-0.2.0/deepinv/datasets/patch_dataset.py

@@ -0,0 +1,29 @@
+from torch.utils import data
+
+
+class PatchDataset(data.Dataset):
+    r"""
+    Builds the dataset of all patches from a tensor of images.
+
+    :param torch.Tensor imgs: Tensor of images, size: batch size x channels x height x width
+    :param int patch_size: size of patches
+    :param callable: data augmentation. callable object, None for no augmentation.
+    """
+    def __init__(self, imgs, patch_size=6, transforms=None):
+        self.imgs = imgs
+        self.patch_size = patch_size
+        self.patches_per_image = (self.imgs.shape[2]-patch_size+1)*(self.imgs.shape[3]-patch_size+1)
+        self.transforms = transforms
+
+    def __len__(self):
+        return self.imgs.shape[0]*self.patches_per_image
+
+    def __getitem__(self, idx):
+        idx_img = idx // self.patches_per_image
+        idx_in_img = idx % self.patches_per_image
+        idx_x = idx_in_img // (self.imgs.shape[3]-self.patch_size+1)
+        idx_y = idx_in_img % (self.imgs.shape[3]-self.patch_size+1)
+        patch = self.imgs[idx_img, :, idx_x:idx_x+self.patch_size, idx_y:idx_y+self.patch_size]
+        if self.transforms and False:
+            patch = self.transforms(patch)
+        return patch.reshape(-1), idx

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/loss/__init__.py

@@ -4,7 +4,9 @@ from deepinv.loss.moi import MOILoss
 from deepinv.loss.sup import SupLoss
 from deepinv.loss.score import ScoreLoss
 from deepinv.loss.tv import TVLoss
+from deepinv.loss.r2r import R2RLoss
 from deepinv.loss.sure import SureGaussianLoss, SurePoissonLoss, SurePGLoss
 from deepinv.loss.regularisers import JacobianSpectralNorm, FNEJacobianSpectralNorm
 from deepinv.loss.measplit import SplittingLoss, Neighbor2Neighbor
-from deepinv.loss.metric import LpNorm, CharbonnierLoss
+from deepinv.loss.metric import LpNorm, PSNR, SSIM, LPIPS, NIQE
+from deepinv.loss.loss import Loss

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/loss/ei.py

@@ -1,8 +1,9 @@
 import torch
 import torch.nn as nn
+from deepinv.loss.loss import Loss
 
 
-class EILoss(nn.Module):
+class EILoss(Loss):
     r"""
     Equivariant imaging self-supervised loss.
 

deepinv-0.2.0/deepinv/loss/loss.py

@@ -0,0 +1,29 @@
+import torch
+
+
+class Loss(torch.nn.Module):
+    r"""
+    Base class for all loss/metric functions.
+
+    Sets a template for the loss functions, whose forward method must follow the input parameters in
+    :meth:`deepinv.loss.Loss.forward`.
+    """
+
+    def __init__(self):
+        super(Loss, self).__init__()
+
+    def forward(self, x_net, x, y, physics, model, **kwargs):
+        r"""
+        Computes the loss.
+
+        :param torch.Tensor x_net: Reconstructed image :math:`\inverse{y}`.
+        :param torch.Tensor x: Reference image.
+        :param torch.Tensor y: Measurement.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+        :param torch.nn.Module model: Reconstruction function.
+
+        :return: (torch.Tensor) loss, the tensor size might be (1,) or (batch size,).
+        """
+        raise NotImplementedError(
+            "The method 'forward' must be implemented in the subclass."
+        )

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/loss/mc.py

@@ -1,8 +1,8 @@
 import torch
-import torch.nn as nn
+from deepinv.loss.loss import Loss
 
 
-class MCLoss(nn.Module):
+class MCLoss(Loss):
     r"""
     Measurement consistency loss
 

{deepinv-0.1.0.dev0 → deepinv-0.2.0}/deepinv/loss/measplit.py

@@ -1,9 +1,10 @@
 import torch
 from deepinv.physics import Inpainting
 import numpy as np
+from deepinv.loss.loss import Loss
 
 
-class SplittingLoss(torch.nn.Module):
+class SplittingLoss(Loss):
     r"""
     Measurement splitting loss.
 
@@ -13,7 +14,7 @@ class SplittingLoss(torch.nn.Module):
 
     .. math::
 
-        \frac{m}{m_2}\| y_2 - A_2 \inversef{y_1,A_1}\|^2
+        \frac{m}{m_2}\| y_2 - A_2 \inversef{y_1}{A_1}\|^2
 
     where :math:`R` is the trainable network. See https://pubmed.ncbi.nlm.nih.gov/32614100/.
 
@@ -55,8 +56,8 @@ class SplittingLoss(torch.nn.Module):
             mask[..., start::stride, start::stride] = 0.0
 
         # create inpainting masks
-        inp = Inpainting(tsize, mask)
-        inp2 = Inpainting(tsize, 1 - mask)
+        inp = Inpainting(tsize, mask, device=y.device)
+        inp2 = Inpainting(tsize, 1 - mask, device=y.device)
 
         # concatenate operators
         physics1 = inp * physics  # A_1 = P*A
@@ -72,7 +73,7 @@ class SplittingLoss(torch.nn.Module):
         return loss_ms
 
 
-class Neighbor2Neighbor(torch.nn.Module):
+class Neighbor2Neighbor(Loss):
     r"""
     Neighbor2Neighbor loss.