# This file is part of LION library
# License : GPL-3
#
# Author : Zakhar Shumaylov, Charlie Shoebridge
# Modifications: Ander Biguri
# =============================================================================
from __future__ import annotations
import numpy as np
from typing import Optional
import torch
from torch.optim.optimizer import Optimizer
from LION.models.LIONmodel import LIONmodel, ModelInputType
from LION.optimizers.LIONsolver import LIONsolver
from LION.utils.parameter import LIONParameter
from LION.utils.citation import citation
from LION.utils.parameter import OperatorParameters
[docs]
class ARSolver(LIONsolver):
"""Trains an adversarial regulariser (Lunz et al., 2018).
Training compares clean images with reconstructions of the measurements
(``recon_fn``, by default the operator's: FDK for CT) (WGAN loss with
gradient penalty). The WGAN loss doesn't measure reconstructions, so
validation and testing reconstruct with a ``LearnedRegularizerReconstructor``
you make and give (``set_validation(..., reconstructor=...)``): the
solver uses it with the regulariser being trained. Its settings (e.g.
``alpha``, which ``estimate_alpha`` estimates as in the paper) are the
reconstructor's. To reconstruct with a trained regulariser, use the same
class with the model.
"""
def __init__(
self,
model: LIONmodel,
optimizer: Optimizer,
operator_params: OperatorParameters = None,
solver_params: LIONParameter | None = None,
model_regularization=None,
scheduler=None,
scheduler_step: str = "epoch",
verbose: int = 1,
device: torch.device = None,
seed: Optional[int] = 0,
deterministic: bool = False,
) -> None:
super().__init__(
model,
optimizer,
self.wgan_loss,
operator_params=operator_params,
solver_params=solver_params,
model_regularization=model_regularization,
scheduler=scheduler,
scheduler_step=scheduler_step,
verbose=verbose,
device=device,
seed=seed,
deterministic=deterministic,
)
# Its models reconstruct only through a reconstructor (LION.evaluation
# refuses their files alone).
reconstructed_with = "LearnedRegularizerReconstructor"
[docs]
@staticmethod
def default_parameters() -> LIONParameter:
"""``lambd_gp``: the weight of the gradient penalty in the WGAN
training loss (0.1). ``recon_fn``: how measurements are reconstructed
into the "bad" reconstructions the regulariser learns to tell from
clean images, ``recon_fn(y, op)``; "default": the operator's default
reconstruction (FDK for CT), replaced by that function when the
solver is made."""
params = LIONParameter()
params.lambd_gp = 1e-1
params.recon_fn = "default"
return params
[docs]
def set_validation(
self,
validation_loader,
validation_freq,
validation_fn=None,
validation_fname=None,
reconstructor=None,
):
"""As ``LIONsolver.set_validation``; ``reconstructor`` (a
``LearnedRegularizerReconstructor``) is needed. By default, validate
with the MSE between reconstruction and target."""
_needs(reconstructor, "set_validation")
if validation_fn is None:
validation_fn = torch.nn.MSELoss()
super().set_validation(
validation_loader,
validation_freq,
validation_fn,
validation_fname,
reconstructor,
)
[docs]
def set_testing(self, test_loader, testing_fn=None, reconstructor=None):
"""As ``LIONsolver.set_testing``; ``reconstructor`` is needed."""
_needs(reconstructor, "set_testing")
if testing_fn is None:
testing_fn = torch.nn.MSELoss()
super().set_testing(test_loader, testing_fn, reconstructor)
[docs]
def wgan_loss(self, y, target_batch):
"""The WGAN loss with gradient penalty: the regulariser's mean on
clean images (``target_batch``) minus its mean on reconstructions of
``y``, plus ``lambd_gp`` times the gradient penalty at random points
between them."""
bad_recon = self.reconstruct_measurements(y, self.op)
if self.do_normalise and self.model.get_input_type() == ModelInputType.IMAGE:
bad_recon = self.model.normalise(bad_recon)
target_batch = self.model.normalise(target_batch)
epsilon = torch.Tensor(
np.random.random((target_batch.shape[0], 1, 1, 1))
).type_as(target_batch)
interpolates = (
epsilon * target_batch + ((1 - epsilon) * bad_recon)
).requires_grad_(True)
net_interpolates = self.model(interpolates)
fake = (
torch.Tensor(net_interpolates.shape)
.fill_(1.0)
.type_as(target_batch)
.requires_grad_(False)
)
gradients = torch.autograd.grad(
outputs=net_interpolates,
inputs=interpolates,
grad_outputs=fake,
create_graph=True,
retain_graph=True,
only_inputs=True,
)[0]
gradients = gradients.view(gradients.shape[0], -1)
loss = (
(self.model(target_batch)).mean()
- (self.model(bad_recon)).mean()
+ self.solver_params.lambd_gp
* (((gradients.norm(2, dim=1) - 1)) ** 2).mean()
)
return loss
[docs]
def mini_batch_step(self, y, target_batch) -> float:
"""The WGAN loss of the batch (``wgan_loss``)."""
return self._regularized(self.wgan_loss(y, target_batch))
[docs]
@staticmethod
def cite(cite_format="bib"):
"""Where ARSolver comes from: its citation, BibTeX ("bib") or "MLA"."""
return citation(
cite_format,
bib="""
@article{lunz2018adversarial,
title={Adversarial regularizers in inverse problems},
author={Lunz, Sebastian and {\\"O}ktem, Ozan and Sch{\\"o}nlieb, Carola-Bibiane},
journal={Advances in neural information processing systems},
volume={31},
year={2018}
}
""",
mla="""
Lunz, Sebastian, Ozan Öktem, and Carola-Bibiane Schönlieb
"Adversarial regularizers in inverse problems."
neural information processing systems 31 (2018).
arXiv:2008.02839 (2020).
""",
)
def _needs(reconstructor, where):
if reconstructor is None:
raise ValueError(
f"ARSolver.{where} needs a reconstructor: the regulariser alone "
"doesn't reconstruct. Give e.g. "
"reconstructor=LearnedRegularizerReconstructor(operator, "
"parameters=...) (alpha: estimate_alpha(op, loader))"
)