Module flowcon.CNF.neural_odes.diffeq_layers.container

Classes

class MixtureODELayer (experts)

Produces a mixture of experts where output = sigma(t) * f(t, x). Time-dependent weights sigma(t) help learn to blend the experts without resorting to a highly stiff f. Supports both regular and diffeq experts.

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

Expand source code

class MixtureODELayer(nn.Module):
    """Produces a mixture of experts where output = sigma(t) * f(t, x).
    Time-dependent weights sigma(t) help learn to blend the experts without resorting to a highly stiff f.
    Supports both regular and diffeq experts.
    """

    def __init__(self, experts):
        super(MixtureODELayer, self).__init__()
        assert len(experts) > 1
        wrapped_experts = [diffeq_wrapper(ex) for ex in experts]
        self.experts = nn.ModuleList(wrapped_experts)
        self.mixture_weights = nn.Linear(1, len(self.experts))

    def forward(self, t, y):
        dys = []
        for f in self.experts:
            dys.append(f(t, y))
        dys = torch.stack(dys, 0)
        weights = self.mixture_weights(t).view(-1, *([1] * (dys.ndimension() - 1)))

        dy = torch.sum(dys * weights, dim=0, keepdim=False)
        return dy

Ancestors

• torch.nn.modules.module.Module

Class variables

var call_super_init : bool

var dump_patches : bool

var training : bool

Methods

def forward(self, t, y) ‑> Callable[..., Any]

Define 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 :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class SequentialDiffEq (*layers)

A container for a sequential chain of layers. Supports both regular and diffeq layers.

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

Expand source code

class SequentialDiffEq(nn.Module):
    """A container for a sequential chain of layers. Supports both regular and diffeq layers.
    """

    def __init__(self, *layers):
        super(SequentialDiffEq, self).__init__()
        self.layers = nn.ModuleList([diffeq_wrapper(layer) for layer in layers])

    def forward(self, t, x):
        for layer in self.layers:
            x = layer(t, x)
        return x

Ancestors

• torch.nn.modules.module.Module

Subclasses

• ResNet

Class variables

var call_super_init : bool

var dump_patches : bool

var training : bool

Methods

def forward(self, t, x) ‑> Callable[..., Any]

Define 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 :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.