LION.models.torch_model

Classes

TorchModel(module_class[, arguments, ...])

A plain torch module as a LIONmodel: module_class(**arguments).

Functions

noise_level_input(module_class)

How instances of module_class take the noise level: from NOISE_LEVEL, else a keyword of NOISE_LEVEL_KEYWORDS that its forward has; None if neither (it takes none, as far as LION knows).

---

class LION.models.torch_model.TorchModel(module_class, arguments=None, model_parameters=None, noise_level=None)

Bases: LIONmodel

A plain torch module as a LIONmodel: module_class(**arguments).

The class and its arguments are recorded, so the model can be saved, loaded and made again like any LIONmodel, and used where LION needs one (e.g. as GSD’s denoiser):

model = TorchModel(MyNet, {"channels": 32})
model.save("my_net")
model = TorchModel.load("my_net", allow=["my_package."])

module_class must be importable where the model is loaded: a class from torch.nn or LION loads as it is; one from your own module needs allow=["your_module."] when loading (one defined in the script that loads it: allow=["__main__."]). arguments must be values a spec can hold (numbers, strings, lists, dicts, LION parameters, …), not tensors.

model_parameters: LION’s own (model_input_type, IMAGE by default).

noise_level: how the module takes the noise level, for a denoiser: “channel” (an extra input channel, constant at the noise level) or the name of its keyword argument (e.g. “sigma”). None: as noise_level_input finds (NOISE_LEVEL, then the module’s forward). The model is called as every LION model taking a noise level is, model(x, noise_level=...), and gives it to the module that way. Other arguments go to the module as they are.

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

static default_parameters()

LION’s own parameters: model_input_type IMAGE. The module’s are its arguments.

takes_noise_level()

Whether the module takes the noise level (noise_level, or as found by noise_level_input).

forward(x, *args, noise_level=None, **kwargs)

The module on x, given noise_level the way it takes it (as a channel or a keyword), if it takes one; other arguments as they are.

LION.models.torch_model.noise_level_input(module_class)

How instances of module_class take the noise level: from NOISE_LEVEL, else a keyword of NOISE_LEVEL_KEYWORDS that its forward has; None if neither (it takes none, as far as LION knows).