=============
Customization
=============

Customizing NeuralNet
---------------------

Apart from the :class:`.NeuralNet` base class, we provide
:class:`.NeuralNetClassifier`, :class:`.NeuralNetBinaryClassifier`,
and :class:`.NeuralNetRegressor` for typical classification, binary
classification, and regressions tasks. They should work as drop-in
replacements for sklearn classifiers and regressors.

The :class:`.NeuralNet` class is a little less opinionated about the
incoming data, e.g. it does not determine a loss function by default.
Therefore, if you want to write your own subclass for a special use
case, you would typically subclass from :class:`.NeuralNet`. The
:func:`~skorch.net.NeuralNet.predict` method returns the same output
as :func:`~skorch.net.NeuralNet.predict_proba` by default, which is
the module output (or the first module output, in case it returns
multiple values).

:class:`.NeuralNet` and its subclasses are already very flexible as they are and
should cover many use cases by adjusting the provided parameters or by using
callbacks. However, this may not always be sufficient for your use cases. If you
thus find yourself wanting to customize :class:`.NeuralNet`, please follow the
guidelines in this section.

Methods starting with get_*
^^^^^^^^^^^^^^^^^^^^^^^^^^^

The net provides a few ``get_*`` methods, most notably
:func:`~skorch.net.NeuralNet.get_loss`,
:func:`~skorch.net.NeuralNet.get_dataset`, and
:func:`~skorch.net.NeuralNet.get_iterator`. The intent of these
methods should be pretty self-explanatory, and if you are still not
quite sure, consult their documentations. In general, these methods
are fairly safe to override as long as you make sure to conform to the
same signature as the original.

A short example should serve to illustrate this.
:func:`~skorch.net.NeuralNet.get_loss` is called when the loss is determined.
Below we show an example of overriding :func:`~skorch.net.NeuralNet.get_loss` to
add L1 regularization to our total loss:

.. code:: python

    class RegularizedNet(NeuralNet):
        def __init__(self, *args, lambda1=0.01, **kwargs):
            super().__init__(*args, **kwargs)
            self.lambda1 = lambda1

        def get_loss(self, y_pred, y_true, X=None, training=False):
            loss = super().get_loss(y_pred, y_true, X=X, training=training)
            loss += self.lambda1 * sum([w.abs().sum() for w in self.module_.parameters()])
            return loss

.. note:: This example also regularizes the biases, which you typically
    don't need to do.

It is often a good idea to call ``super`` of the method you override, to make
sure that everything that needs to happen inside that method does happen. If you
don't, you should make sure to take care of everything that needs to happen by
following the original implementation.

Training and validation
^^^^^^^^^^^^^^^^^^^^^^^

If you would like to customize training and validation, there are
several possibilities. Below are the methods that you most likely want
to customize:

The method :func:`~skorch.net.NeuralNet.train_step_single` performs a
single training step. It accepts the current batch of data as input
(as well as the ``fit_params``) and should return a dictionary
containing the ``loss`` and the prediction ``y_pred``. E.g. you should
override this if your dataset returns some non-standard data that
needs custom handling, and/or if your module has to be called in a
very specific way. If you want to, you can still make use of
:func:`~skorch.net.NeuralNet.infer` and
:func:`~skorch.net.NeuralNet.get_loss` but it's not strictly
necessary. Don't call the optimizer in this method, this is handled by
the next method.

The method :func:`~skorch.net.NeuralNet.train_step` defines the
complete training procedure performed for each batch. It accepts the
same arguments as :func:`~skorch.net.NeuralNet.train_step_single` but
it differs in that it defines the training closure passed to the
optimizer, which for instance could be called more than once (e.g. in
case of L-BFGS). You might override this if you deal with non-standard
training procedures, as e.g. gradient accumulation.

The method :func:`~skorch.net.NeuralNet.validation_step` is
responsible for calculating the prediction and loss on the validation
data (remember that skorch uses an internal validation set for
reporting, early stopping, etc.). Similar to
:func:`~skorch.net.NeuralNet.train_step_single`, it receives the batch
and ``fit_params`` as input and should return a dictionary containing
``loss`` and ``y_pred``. Most likely, when you need to customize
:func:`~skorch.net.NeuralNet.train_step_single`, you'll need to
customize :func:`~skorch.net.NeuralNet.validation_step` accordingly.

Finally, the method :func:`~skorch.net.NeuralNet.evaluation_step` is
called to you use the net for inference, e.g. when calling
:func:`~skorch.net.NeuralNet.forward` or
:func:`~skorch.net.NeuralNet.predict`. You may want to modify this if,
e.g., you want your model to behave differently during training and
during prediction.

You should also be aware that some methods are better left
untouched. E.g., in most cases, the following methods should *not* be
overridden:

* :func:`~skorch.net.NeuralNet.fit`
* :func:`~skorch.net.NeuralNet.partial_fit`
* :func:`~skorch.net.NeuralNet.fit_loop`
* :func:`~skorch.net.NeuralNet.run_single_epoch`

The reason why these methods should stay untouched is because they
perform some book keeping, like making sure that callbacks are handled
or writing logs to the ``history``. If you do need to override these,
make sure that you perform the same book keeping as the original
methods.

Initialization and custom modules
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

The method :func:`~skorch.net.NeuralNet.initialize` is responsible for
initializing all the components needed by the net, e.g. the module and
the optimizer. For this, it calls specific initialization methods,
such as :func:`~skorch.net.NeuralNet.initialize_module` and
:func:`~skorch.net.NeuralNet.initialize_optimizer`. If you'd like to
customize the initialization behavior, you should override the
corresponding methods. Following sklearn conventions, the created
components should be set as an attribute with a trailing underscore as
the name, e.g. ``module_`` for the initialized module.

A possible modification you may want to make is to add more modules, criteria,
and optimizers to your net. This is possible in skorch by following the
guidelines below. If you do this, your custom modules and optimizers will be
treated as "first class citizens" in skorch land. This means:

1. The parameters of your custom modules are automatically passed to the
   optimizer (but you can modify this behavior).
2. skorch takes care of moving your modules to the correct device.
3. skorch takes care of setting the training/eval mode correctly.
4. When a module needs to be re-initialized because ``set_params`` was called,
   all modules and optimizers that may depend on it are also re-initialized.
   This is for instance important for the optimizer, which must know about the
   parameters of the newly initialized module.
5. You can pass arguments to the custom modules and optimizers using the now
   familiar double-underscore notation. E.g., you can initialize your net like
   this:

.. code:: python

    net = MyNet(
        module=MyModule,
        module__num_units=100,

        othermodule=MyOtherModule,
        othermodule__num_units=200,
    )
    net.fit(X, y)

A word about the distinction between modules and criteria made by skorch:
Typically, criteria are also just subclasses of PyTorch
:class:`~torch.nn.Module`. As such, skorch moves them to CUDA if that is the
indicated device and will even pass parameters of criteria to the optimizers, if
there are any. This can be useful when e.g. training GANs, where you might
implement the discriminator as the criterion (and the generator as the module).

A difference between module and criterion is that the output of modules are used
for generating the predictions and are thus returned by
:func:`~skorch.net.NeuralNet.predict` etc. In contrast, the output of the
criterion is used for calculating the loss and should therefore be a scalar.

skorch assumes that criteria may depend on the modules. Therefore, if a module
is re-initialized, all criteria are also re-initialized, but not vice-versa. On
top of that, the optimizer is re-initialized when either modules or criteria
are changed.

So after all this talk, what are the aforementioned guidelines to add your own
modules, criteria, and optimizers? You have to follow these rules:

1. Initialize them during their respective ``initialize_`` methods, e.g. modules
   should be set inside :func:`~skorch.net.NeuralNet.initialize_module`.
2. If they have learnable parameters, they should be instances of
   :class:`~torch.nn.Module`. Optimizers should be instances of
   :class:`~torch.optim.Optimizer`.
3. Their names should end on an underscore. This is true for all attributes that
   are created during ``initialize`` and distinguishes them from arguments
   passed to ``__init__``. So a name for a custom module could be ``mymodule_``.
4. Inside the initialization method, use
   :meth:`skorch.net.NeuralNet.get_params_for` (or, if dealing with an
   optimizer, :meth:`skorch.net.NeuralNet.get_params_for_optimizer`) to retrieve
   the arguments for the constructor of the instance.

Here is an example of how this could look like in practice:

.. code:: python

    class MyNet(NeuralNet):
        def initialize_module(self):
            super().initialize_module()

            # add an additional module called 'module2_'
            params = self.get_params_for('module2')
            self.module2_ = Module2(**params)
            return self

        def initialize_criterion(self):
            super().initialize_criterion()

            # add an additional criterion called 'other_criterion_'
            params = self.get_params_for('other_criterion')
            self.other_criterion_ = nn.BCELoss(**params)
            return self

        def initialize_optimizer(self):
            # first initialize the normal optimizer
            named_params = self.module_.named_parameters()
            args, kwargs = self.get_params_for_optimizer('optimizer', named_params)
            self.optimizer_ = self.optimizer(*args, **kwargs)

            # next add an another optimizer called 'optimizer2_' that is
            # only responsible for training 'module2_'
            named_params = self.module2_.named_parameters()
            args, kwargs = self.get_params_for_optimizer('optimizer2', named_params)
            self.optimizer2_ = torch.optim.SGD(*args, **kwargs)
            return self

        ...  # additional changes


    net = MyNet(
        ...,
        module2__num_units=123,
        other_criterion__reduction='sum',
        optimizer2__lr=0.1,
    )
    net.fit(X, y)

    # set_params works
    net.set_params(optimizer2__lr=0.05)
    net.partial_fit(X, y)

    # grid search et al. works
    search = GridSearchCV(net, {'module2__num_units': [10, 50, 100]}, ...)
    search.fit(X, y)

In this example, a new criterion, a new module, and a new optimizer
were added. Of course, additional changes should be made to the net so
that those new components are actually being used for something, but
this example should illustrate how to start. Since the rules outlined
above are being followed, we can use grid search on our customly
defined components.

.. note:: In the example above, the parameters of ``module_`` are trained by
          ``optimzer_`` and the parameters of ``module2_`` are trained by
          ``optimizer2_``. To conveniently obtain the parameters of all modules,
          call the method :func:`~skorch.net.NeuralNet.get_all_learnable_params`.