Plug-and-Play Regularization

A group of regularization terms that can not be directly written down as function are learned plug-and-play (PnP) priors. These are terms based on deep neural networks, which are trainted to implement the proximal map corresponding to the regularization term. Such a PnP prior can be used in the same way as any other regularization term.

The following example shows how to use a PnP prior in the context of the Kaczmarz solver.

using RegularizedLeastSquares
A = randn(32, 16)
x = randn(16)
b = A*x;

For the documentation we will just use the identity function as a placeholder for the PnP prior.

model = identity
identity (generic function with 1 method)

In practice, you would replace this with a neural network:

using Flux
model = Flux.loadmodel!(model, ...)

The model can then be used together with the PnPRegularization term:

reg = PnPRegularization(1.0; model = model, shape = [16]);

Since models often expect a specific input range, we can use the MinMaxTransform to normalize the input:

reg = PnPRegularization(1.0; model = model, shape = [16], input_transform = RegularizedLeastSquares.MinMaxTransform);

Custom input transforms can be implemented by passing something callable as the input_transform keyword argument. For more details see the PnPRegularization documentation.

The regularization term can then be used in the solver:

solver = createLinearSolver(Kaczmarz, A; reg = reg, iterations = 32)
x_approx = solve!(solver, b)
16-element Vector{Float64}:
 -0.016645515429894342
 -0.7015139729646289
  0.23577353393146683
 -1.3780583440254406
 -0.6769997880493167
 -0.22614659153825412
  0.001475371713707041
  1.4403403763322167
 -1.2911069997038678
 -0.9155552170186194
  0.5557743078048285
 -0.46487512690783517
 -0.7529450586773503
 -0.08319171301213335
  0.5117173508705821
 -0.8895085103780653

This page was generated using Literate.jl.