Skip to contents

The setup function is used to set important attributes and method for nn_modules to be used with luz.

Usage

setup(module, loss = NULL, optimizer = NULL, metrics = NULL, backward = NULL)

Arguments

module

(nn_module) The nn_module that you want set up.

loss

(function, optional) An optional function with the signature function(input, target). It's only requires if your nn_module doesn't implement a method called loss.

optimizer

(torch_optimizer, optional) A function with the signature function(parameters, ...) that is used to initialize an optimizer given the model parameters.

metrics

(list, optional) A list of metrics to be tracked during the training procedure.

backward

(function) A functions that takes the loss scalar values as it's parameter. It must call $backward() or torch::autograd_backward(). In general you don't need to set this parameter unless you need to customize how luz calls the backward(), for example, if you need to add additional arguments to the backward call. Note that this becomes a method of the nn_module thus can be used by your custom step() if you override it.

Value

A luz module that can be trained with fit().

Details

It makes sure the module have all the necessary ingredients in order to be fitted.