torchdyn.models

Continuous or Implicit neural model templates. Pretrained models (eventually).

Submodules

torchdyn.models.cnf module

class torchdyn.models.cnf.CNF(net, trace_estimator=None, noise_dist=None, order=1)

Continuous Normalizing Flow

Parameters

• net (nn.Module) – function parametrizing the datasets vector field.

• trace_estimator (Callable) – specifies the strategy to otbain Jacobian traces. Options: (autograd_trace, hutch_trace)

• noise_dist (torch.distributions.Distribution) – distribution of noise vectors sampled for stochastic trace estimators. Needs to have a .sample method.

• order (int) – specifies parameters of the Neural DE.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

torchdyn.models.cnf.autograd_trace(x_out, x_in, **kwargs)

Standard brute-force means of obtaining trace of the Jacobian, O(d) calls to autograd

torchdyn.models.cnf.hutch_trace(x_out, x_in, noise=None, **kwargs)

Hutchinson’s trace Jacobian estimator, O(1) call to autograd

torchdyn.models.energy module

class torchdyn.models.energy.ConservativeLinearSNF(energy, J)

Stable Neural Flows: https://arxiv.org/abs/2003.08063 A generalization of Hamiltonian Neural Networks and other energy-based parametrization of Neural ODEs Conservative version with energy preservation. Input assumed to be of dimensions batch, dim

Parameters

• energy – function parametrizing the energy.

• J – network parametrizing the skew-symmetric interconnection matrix

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchdyn.models.energy.GNF(energy)

Gradient Neural Flows version of SNFs: https://arxiv.org/abs/2003.08063 :param energy: function parametrizing the energy. :type energy: nn.Module

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchdyn.models.energy.HNN(net)

Hamiltonian Neural ODE

Parameters

net (nn.Module) – function parametrizing the vector field.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchdyn.models.energy.LNN(net)

Lagrangian Neural Network.

Parameters

net (nn.Module) –

Notes

LNNs are currently quite slow. Improvements will be made whenever functorch is either merged upstream or included as a dependency.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.