deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/utils/nn.py

@@ -0,0 +1,213 @@
+import torch
+import os
+import numpy as np
+
+
+class TensorList:
+    r"""
+
+    Represents a list of :class:`torch.Tensor` with different shapes.
+    It allows to sum, flatten, append, etc. lists of tensors seamlessly, in a
+    similar fashion to :class:`torch.Tensor`.
+
+    :param x: a list of :class:`torch.Tensor`, a single :class:`torch.Tensor` or a TensorList.
+    """
+
+    def __init__(self, x):
+        super().__init__()
+
+        if isinstance(x, list) or isinstance(x, TensorList):
+            self.x = list(x)
+        elif isinstance(x, torch.Tensor):
+            self.x = [x]
+        else:
+            raise TypeError("x must be a list of torch.Tensor or a single torch.Tensor")
+
+        self.shape = [xi.shape for xi in self.x]
+
+    def __len__(self):
+        r"""
+        Returns the number of tensors in the list.
+        """
+        return len(self.x)
+
+    def __getitem__(self, item):
+        r"""
+        Returns the ith tensor in the list.
+        """
+        return self.x[item]
+
+    def flatten(self):
+        r"""
+        Returns a :class:`torch.Tensor` with a flattened version of the list of tensors.
+        """
+        return torch.cat([xi.flatten() for xi in self.x])
+
+    def append(self, other):
+        r"""
+        Appends a :class:`torch.Tensor` or a list of :class:`torch.Tensor` to the list.
+
+        """
+        if isinstance(other, list):
+            self.x += other
+        elif isinstance(other, TensorList):
+            self.x += other.x
+        elif isinstance(other, torch.Tensor):
+            self.x.append(other)
+        else:
+            raise TypeError(
+                "the appended item must be a list of :class:`torch.Tensor` or a single :class:`torch.Tensor`"
+            )
+        return self
+
+    def __add__(self, other):
+        r"""
+
+        Adds two TensorLists. The sizes of the tensor lists must match.
+
+        """
+        if not isinstance(other, list) and not isinstance(other, TensorList):
+            return TensorList([xi + other for xi in self.x])
+        else:
+            return TensorList([xi + otheri for xi, otheri in zip(self.x, other)])
+
+    def __mul__(self, other):
+        r"""
+
+        Multiply two TensorLists. The sizes of the tensor lists must match.
+
+        """
+        if not isinstance(other, list) and not isinstance(other, TensorList):
+            return TensorList([xi * other for xi in self.x])
+        else:
+            return TensorList([xi * otheri for xi, otheri in zip(self.x, other)])
+
+    def __truediv__(self, other):
+        r"""
+
+        Divide two TensorLists. The sizes of the tensor lists must match.
+
+        """
+        if not isinstance(other, list) and not isinstance(other, TensorList):
+            return TensorList([xi / other for xi in self.x])
+        else:
+            return TensorList([xi / otheri for xi, otheri in zip(self.x, other)])
+
+    def __neg__(self):
+        r"""
+
+        Negate a TensorList.
+        """
+        return TensorList([-xi for xi in self.x])
+
+    def __sub__(self, other):
+        r"""
+
+        Substract two TensorLists. The sizes of the tensor lists must match.
+
+        """
+        if not isinstance(other, list) and not isinstance(other, TensorList):
+            return TensorList([xi - other for xi in self.x])
+        else:
+            return TensorList([xi - otheri for xi, otheri in zip(self.x, other)])
+
+
+def randn_like(x):
+    r"""
+    Returns a :class:`deepinv.utils.TensorList` or :class:`torch.Tensor`
+    with the same type as x, filled with standard gaussian numbers.
+    """
+    if isinstance(x, torch.Tensor):
+        return torch.randn_like(x)
+    else:
+        return TensorList([torch.randn_like(xi) for xi in x])
+
+
+def rand_like(x):
+    r"""
+    Returns a :class:`deepinv.utils.TensorList` or :class:`torch.Tensor`
+    with the same type as x, filled with random uniform numbers in [0,1].
+    """
+    if isinstance(x, torch.Tensor):
+        return torch.rand_like(x)
+    else:
+        return TensorList([torch.rand_like(xi) for xi in x])
+
+
+def zeros_like(x):
+    r"""
+    Returns a :class:`deepinv.utils.TensorList` or :class:`torch.Tensor`
+    with the same type as x, filled with zeros.
+    """
+    if isinstance(x, torch.Tensor):
+        return torch.zeros_like(x)
+    else:
+        return TensorList([torch.zeros_like(xi) for xi in x])
+
+
+def ones_like(x):
+    r"""
+    Returns a :class:`deepinv.utils.TensorList` or :class:`torch.Tensor`
+    with the same type as x, filled with ones.
+    """
+    if isinstance(x, torch.Tensor):
+        return torch.ones_like(x)
+    else:
+        return TensorList([torch.ones_like(xi) for xi in x])
+
+
+def get_freer_gpu():
+    """
+    Returns the GPU device with the most free memory.
+
+    """
+    try:
+        if os.name == "posix":
+            os.system("nvidia-smi -q -d Memory |grep -A5 GPU|grep Free >tmp")
+            memory_available = [int(x.split()[2]) for x in open("tmp", "r").readlines()]
+        else:
+            os.system('bash -c "nvidia-smi -q -d Memory |grep -A5 GPU|grep Free >tmp"')
+            memory_available = [int(x.split()[2]) for x in open("tmp", "r").readlines()]
+        idx = np.argmax(memory_available)
+        device = torch.device(f"cuda:{idx}")
+        print(f"Selected GPU {idx} with {np.max(memory_available)} MB free memory ")
+    except:
+        device = torch.device(f"cuda")
+        print("Couldn't find free GPU")
+
+    return device
+
+
+def save_model(
+    epoch, model, optimizer, ckp_interval, epochs, loss, save_path, eval_psnr=None
+):
+    if (epoch > 0 and epoch % ckp_interval == 0) or epoch + 1 == epochs:
+        os.makedirs(save_path, exist_ok=True)
+
+        state = {
+            "epoch": epoch,
+            "state_dict": model.state_dict(),
+            "loss": loss,
+            "optimizer": optimizer.state_dict(),
+        }
+        if eval_psnr is not None:
+            state["eval_psnr"] = eval_psnr
+        torch.save(state, os.path.join(save_path, "ckp_{}.pth.tar".format(epoch)))
+    pass
+
+
+def load_checkpoint(model, path_checkpoint, device):
+    checkpoint = torch.load(path_checkpoint, map_location=device)
+    model.load_state_dict(checkpoint["state_dict"])
+    return model
+
+
+def investigate_model(model, idx_max=1, check_name="iterator.g_step.g_param.0"):
+    for idx, (name, param) in enumerate(model.named_parameters()):
+        if param.requires_grad and (idx < idx_max or check_name in name):
+            print(
+                name,
+                param.data.flatten()[0],
+                "gradient norm = ",
+                param.grad.detach().data.norm(2),
+            )

deepinv/utils/optimization.py

@@ -0,0 +1,108 @@
+import torch.nn as nn
+
+
+class NeuralIteration(nn.Module):
+    def __init__(self):
+        super(NeuralIteration, self).__init__()
+
+    def init(self, backbone_blocks, step_size=1.0, iterations=1):
+        self.iterations = iterations
+        self.n_blocks = len(backbone_blocks) if isinstance(backbone_blocks, list) else 1
+        self.blocks = backbone_blocks
+
+        self.register_parameter(
+            name="step_size",
+            param=torch.nn.Parameter(
+                step_size * torch.ones(self.iterations), requires_grad=True
+            ),
+        )
+        if self.n_blocks > 1:  # weight_tied=False (many blocks)
+            assert (
+                self.n_blocks == iterations
+            ), "'# blocks' does not equal to 'iterations'"
+            self.blocks = torch.nn.ModuleList(
+                [backbone_blocks[_] for _ in range(iterations)]
+            )
+        else:  # weight_tied=True (only one block)
+            self.blocks = torch.nn.ModuleList([backbone_blocks])
+
+    def forward(self, y, physics, x_init=None):
+        return physics.A_adjoint(y)
+
+    @staticmethod
+    def measurement_consistency_grad(physics, x, y):
+        # grad(||y-Ax||) = A^T(y-Ax)
+        return physics.A_adjoint(
+            y.to(physics.device) - physics.A(x.to(physics.device))
+        ).to(x.device)
+
+
+class GradientDescent(NeuralIteration):
+    def __init__(self, backbone_blocks, step_size=1.0, iterations=1):
+        super(GradientDescent, self).__init__()
+        self.name = "gd"
+        self.init(backbone_blocks, step_size, iterations)
+
+    def forward(self, y, physics, x_init=None):
+        x = x_init.clone() if x_init is not None else physics.A_adjoint(y)
+        for t in range(self.iterations):
+            if self.n_blocks == 1:
+                t = 0
+            x = (
+                self.blocks[t](x)
+                + x
+                + self.step_size[t]
+                * NeuralIteration.measurement_consistency_grad(physics, x, y)
+            )
+        return x
+
+
+class ProximalGradientDescent(NeuralIteration):
+    def __init__(self, backbone_blocks, step_size=1.0, iterations=1):
+        super(ProximalGradientDescent, self).__init__()
+        self.name = "pgd"
+        self.init(backbone_blocks, step_size, iterations)
+        self.mc_grad = NeuralIteration.measurement_consistency_grad
+
+    def forward(self, y, physics, x_init=None):
+        x = x_init.clone() if x_init is not None else physics.A_adjoint(y)
+        for t in range(self.iterations):
+            if self.n_blocks == 1:
+                t = 0
+            x = self.blocks[t](
+                x
+                + self.step_size[t]
+                * NeuralIteration.measurement_consistency_grad(physics, x, y)
+            )
+        return x
+
+
+if __name__ == "__main__":
+    import torch
+    import deepinv as dinv
+
+    net = dinv.models.unet().to(dinv.device)
+    physics = dinv.physics.Inpainting([32, 32], device=dinv.device)
+
+    x = torch.randn(10, 1, 32, 32).to(dinv.device)
+    y = physics.A(x)
+    fbp = physics.A_dagger(y)
+    x_rec = net(fbp)
+
+    unroll = ProximalGradientDescent(net, step_size=1.0, iterations=1)
+    x_unroll = unroll(y, physics, x_init=fbp)
+
+    print("iterations=3")
+    iterations = 3
+    step_size = 1.0
+    blocks = [dinv.models.unet().to(dinv.device) for _ in range(iterations)]
+
+    physics = dinv.physics.Inpainting([32, 32], device=dinv.device)
+    x = torch.randn(10, 1, 32, 32).to(dinv.device)
+    y = physics.A(x)
+    fbp = physics.A_dagger(y)
+
+    unroll = ProximalGradientDescent(blocks, step_size=step_size, iterations=iterations)
+    x_unroll = unroll(y, physics, x_init=fbp)
+
+    print(f"iterations={iterations}", x.shape, y.shape, fbp.shape, x_unroll.shape)

deepinv/utils/parameters.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+
+def get_DPIR_params(noise_level_img):
+    r"""
+    Default parameters for the DPIR Plug-and-Play algorithm.
+
+    :param float noise_level_img: Noise level of the input image.
+    """
+    max_iter = 8
+    s1 = 49.0 / 255.0
+    s2 = noise_level_img
+    sigma_denoiser = np.logspace(np.log10(s1), np.log10(s2), max_iter).astype(
+        np.float32
+    )
+    stepsize = (sigma_denoiser / max(0.01, noise_level_img)) ** 2
+    lamb = 1 / 0.23
+    return lamb, list(sigma_denoiser), list(stepsize), max_iter
+
+
+def get_GSPnP_params(problem, noise_level_img):
+    r"""
+    Default parameters for the GSPnP Plug-and-Play algorithm.
+
+    :param str problem: Type of inverse-problem problem to solve. Can be ``deblur``, ``super-resolution``, or ``inpaint``.
+    :param float noise_level_img: Noise level of the input image.
+    """
+    if problem == "deblur":
+        max_iter = 500
+        sigma_denoiser = 1.8 * noise_level_img
+        lamb = 1 / 0.1
+    elif problem == "super-resolution":
+        max_iter = 500
+        sigma_denoiser = 2.0 * noise_level_img
+        lamb = 1 / 0.065
+    elif problem == "inpaint":
+        max_iter = 100
+        sigma_denoiser = 10.0 / 255
+        lamb = 1 / 0.1
+    else:
+        raise ValueError("parameters unknown with this degradation")
+    stepsize = 1.0
+    return lamb, sigma_denoiser, stepsize, max_iter

deepinv/utils/phantoms.py

@@ -0,0 +1,115 @@
+import numpy as np
+import torch
+
+try:
+    import odl
+except:
+    odl = ImportError("The odl package is not installed.")
+
+
+def random_shapes(interior=False):
+    """
+    Generate random shape parameters.
+    Taken from https://github.com/adler-j/adler/blob/master/adler/odl/phantom.py
+    """
+    if interior:
+        x_0 = np.random.rand() - 0.5
+        y_0 = np.random.rand() - 0.5
+    else:
+        x_0 = 2 * np.random.rand() - 1.0
+        y_0 = 2 * np.random.rand() - 1.0
+
+    return (
+        (np.random.rand() - 0.5) * np.random.exponential(0.4),
+        np.random.exponential() * 0.2,
+        np.random.exponential() * 0.2,
+        x_0,
+        y_0,
+        np.random.rand() * 2 * np.pi,
+    )
+
+
+def random_phantom(spc, n_ellipse=50, interior=False):
+    """
+    Generate a random ellipsoid phantom.
+    Taken from https://github.com/adler-j/adler/blob/master/adler/odl/phantom.py
+    """
+    if isinstance(odl, ImportError):
+        raise ImportError(
+            "odl is needed to use generate random phantoms. "
+            "It should be installed with `python3 -m pip install"
+            " https://github.com/odlgroup/odl/archive/master.zip`"
+        ) from odl
+    n = np.random.poisson(n_ellipse)
+    shapes = [random_shapes(interior=interior) for _ in range(n)]
+    return odl.phantom.ellipsoid_phantom(spc, shapes)
+
+
+class RandomPhantomDataset(torch.utils.data.Dataset):
+    """
+    Dataset of random ellipsoid phantoms. The phantoms are generated on the fly.
+    The phantoms are generated using the odl library (https://odlgroup.github.io/odl/).
+
+    :param int size: Size of the phantom (square) image.
+    :param int n_data: Number of phantoms to generate per sample.
+    :param transform: Transformation to apply to the output image.
+    :param float length: Length of the dataset. Useful for iterating the data-loader for a certain nb of iterations.
+    """
+
+    def __init__(self, size=128, n_data=1, transform=None, length=np.inf):
+        self.space = odl.uniform_discr(
+            [-64, -64], [64, 64], [size, size], dtype="float32"
+        )
+        self.transform = transform
+        self.n_data = n_data
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index):
+        """
+        :return tuple : A tuple (phantom, 0) where phantom is a torch tensor of shape (n_data, size, size).
+        """
+        phantom_np = np.array([random_phantom(self.space) for i in range(self.n_data)])
+        phantom = torch.from_numpy(phantom_np).float()
+        if self.transform is not None:
+            phantom = self.transform(phantom)
+        return phantom, 0
+
+
+class SheppLoganDataset(torch.utils.data.Dataset):
+    """
+    Dataset for the single Shepp-Logan phantom. The dataset has length 1.
+    """
+
+    def __init__(self, size=128, n_data=1, transform=None):
+        if isinstance(odl, ImportError):
+            raise ImportError(
+                "odl is needed to use generate the Shepp Logan phantom. "
+                "It should be installed with `python3 -m pip install"
+                " https://github.com/odlgroup/odl/archive/master.zip`"
+            ) from odl
+        self.space = odl.uniform_discr(
+            [-64, -64], [64, 64], [size, size], dtype="float32"
+        )
+        self.transform = transform
+        self.n_data = n_data
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, index):
+        if isinstance(odl, ImportError):
+            raise ImportError(
+                "odl is needed to use generate the Shepp Logan phantom. "
+                "It should be installed with `python3 -m pip install"
+                " https://github.com/odlgroup/odl/archive/master.zip`"
+            ) from odl
+        phantom_np = np.array(
+            [odl.phantom.shepp_logan(self.space, True) for i in range(self.n_data)]
+        )
+        phantom = torch.from_numpy(phantom_np).float()
+        if self.transform is not None:
+            phantom = self.transform(phantom)
+        return phantom, 0