Training the Model
==================

Before training a model, the dataset must be created as described in the `dataset`_ section.

The model can then be trained using a ``ui.json`` interface in `Geoscience Analyst`_. The interface is divided into two tabs: **General** and **Optional** parameters. The **General** tab is shown in :numref:`figure_train_uijson` below.

.. _figure_train_uijson:

.. figure:: ../../images/application/train_main_uijson.png
   :align: center
   :width: 60%
   :alt: train_uijson

   *General tab of the training interface.*

Three options are available in this tab:

 #. **Training dataset**: The ``*.h5`` dataset created in the `dataset`_ section.
 #. **Validation dataset**: The ``*.h5`` dataset created in the `dataset`_ section.
 #. **Output folder**: The folder where the model will be saved.

The **Optional** tab is shown in :numref:`figure_train_optional_uijson` below.

.. _figure_train_optional_uijson:

.. figure:: ../../images/application/train_optional_uijson.png
   :align: center
   :width: 60%
   :alt: train_optional_uijson

   *Optional tab of the training interface.*

The tab is divided into different sections, containing several options to set the training parameters.

**Training options**

 1. **Name**: The name of the model to train.
 2. **Batch size**: The number of samples to use in each batch.
 3. **Epochs**: The number of epochs to train the model.
 4. **Number of samples per epoch**: The number of samples to use in each epoch.
 5. **Samples for training** (*optional*): Randomly select a given number of samples from each dataset. This avoids data imbalance if one layer contains many more tiles than another.

**Learning rate options**

 6. **Learning rate**: The initial learning rate of the optimizer (the `Adam`_ optimizer).
 7. **Factor regularizer**: The regularizer factor for the input block (see the :ref:`wnet_model` section).
 8. **Factor N-cut**: The factor applied to the *N-cut* loss (see the :ref:`wnet_model` section).
 9. **Factor variance**: The factor applied to the variance loss (see the :ref:`wnet_model` section).
 10. **AMSGrad**: If checked, the `AMSGrad`_ optimizer is used.
 11. **Scheduler factor**: The factor applied to the learning rate scheduler. The learning rate scheduler used is the `ReduceLROnPlateau`_ method.
 12. **Scheduler patience**: The patience of the learning rate scheduler.

**N-cut loss options**

 13. **Radius**: The radius for the N-Cut loss (see `soft n-cut loss`_ paper).
 14. **Sigma 1**: The intensity sigma for the N-Cut loss (see `soft n-cut loss`_ paper).
 15. **Sigma 2**: The spatial sigma for the N-Cut loss (see `soft n-cut loss`_ paper).

**Network options**

 16. **Pretrained model** (*optional*): The path to a pretrained model to use.
 17. **Dropout**: The dropout rate to use in the network.
 18. **Output channels**: The number of output channels of the network generator.

**Validation options**

 19. **Validation step**: The number of epochs between each validation.
 20. **Validation patience**: The patience of the validation.
 21. **Samples for validation** (*optional*): Randomly select a given number of samples from each dataset. This avoids data imbalance if one layer contains many more tiles than another.

Once you have set all the parameters, you can start the training by clicking on the ``OK`` button. The training will begin, and the model will be saved in the folder specified in the **General** tab. We highly recommend using a GPU to train the model. During training, at each validation epoch, the loss scores, results of the first validation batch, and the state of the model's weights (as ``.pt`` files) are saved in the **Output folder**.

Training a model is a lengthy process, taking several hours or days depending on the dataset size, the number of epochs, and the computing capabilities of the machine it's trained on. It also requires an iterative trial and error to find the best parameters for the model.

.. _Adam: https://arxiv.org/abs/1412.6980
.. _dataset: prepare_dataset.rst
.. _AMSGrad: https://openreview.net/forum?id=ryQu7f-RZ
.. _ReduceLROnPlateau: https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ReduceLROnPlateau.html
.. _soft n-cut loss: https://arxiv.org/abs/1804.01346
.. _Geoscience Analyst: https://www.mirageoscience.com/mining-industry-software/geoscience-analyst/