LION.models.iterative_unrolled.cLPD

Classes

ContinuousDataProximal(layers, channels[, ...])	cLPD's dual block: a 1x1 convolution in, a neural ODE (first or second order), a 1x1 convolution out.
ContinuousRegProximal(layers, channels[, ...])	cLPD's primal block: as ContinuousDataProximal, on the primal side.
ConvODEFunc(layers, channels[, instance_norm])	The right-hand side of a first-order neural ODE: 3x3 convolutions through the channels listed (with PReLUs, optional instance normalisation).
ConvSODEFunc(layers, channels[, instance_norm])	The right-hand side of a second-order neural ODE: the state is (position, velocity) along the channels; the acceleration is a CNN of both.
InitialVelocity(channels[, instance_norm])	The initial velocity of a second-order neural ODE, from the initial position: convolutions of channels channels.
ODEBlock(odefunc[, tol, adjoint])	Solves the neural ODE odefunc from time 0 to 1 (torchdiffeq).
cLPD(operator_params[, model_parameters])	Learned Primal Dual with continuous blocks (Runkel et al.): LPD whose primal and dual blocks are neural ODEs.

---

class LION.models.iterative_unrolled.cLPD.ContinuousDataProximal(layers, channels, tol=0.001, adjoint=True, do_secon_order=False, solver='rk4', instance_norm=False)

Bases: Module

cLPD’s dual block: a 1x1 convolution in, a neural ODE (first or second order), a 1x1 convolution out.

layers and channels: as dataProximal’s; tol, adjoint and solver: the ODE’s (ODEBlock); do_secon_order: a second-order ODE; instance_norm: in the ODE’s convolutions.

forward(x)

The block’s update of the dual variables, from x.

class LION.models.iterative_unrolled.cLPD.ContinuousRegProximal(layers, channels, tol=0.001, adjoint=False, do_secon_order=False, solver='rk4', instance_norm=False)

Bases: Module

cLPD’s primal block: as ContinuousDataProximal, on the primal side.

As ContinuousDataProximal’s.

forward(x)

The block’s update of the primal variables, from x.

class LION.models.iterative_unrolled.cLPD.ConvODEFunc(layers, channels, instance_norm=False)

Bases: Module

The right-hand side of a first-order neural ODE: 3x3 convolutions through the channels listed (with PReLUs, optional instance normalisation).

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(t, x)

dx/dt at x (t unused: autonomous); counts the evaluations (nfe).

class LION.models.iterative_unrolled.cLPD.ConvSODEFunc(layers, channels, instance_norm=False)

Bases: Module

The right-hand side of a second-order neural ODE: the state is (position, velocity) along the channels; the acceleration is a CNN of both.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(t, x)

The state’s derivative at x: (velocity, acceleration).

class LION.models.iterative_unrolled.cLPD.InitialVelocity(channels, instance_norm=False)

Bases: Module

The initial velocity of a second-order neural ODE, from the initial position: convolutions of channels channels.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x0)

The ODE’s state at time 0: the position x0 and its velocity, along the channels.

class LION.models.iterative_unrolled.cLPD.ODEBlock(odefunc, tol=0.001, adjoint=False)

Bases: Module

Solves the neural ODE odefunc from time 0 to 1 (torchdiffeq).

odefunc: the ODE’s right-hand side; tol: the solver’s relative and absolute tolerance; adjoint: gradients by the adjoint method (less memory) instead of through the solver.

forward(x, eval_times=None, solver='rk4')

The ODE’s solution at time 1 from x (or at each of eval_times), with solver (a torchdiffeq method; rk4 by default).

class LION.models.iterative_unrolled.cLPD.cLPD(operator_params, model_parameters=None)

Bases: LPD

Learned Primal Dual with continuous blocks (Runkel et al.): LPD whose primal and dual blocks are neural ODEs.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

static default_parameters()

LPD’s parameters with 4-layer blocks, plus do_secon_order (second-order ODE blocks; False).

static continous_LPD_paper()

The paper’s cLPD (Runkel et al., 2024): 10 iterations, learned steps, first-order blocks.

static cite(cite_format='bib')

Where cLPD comes from: its citation, BibTeX (“bib”) or “MLA”.