Skip to contents

Base Class for Torch Learners

Source: R/LearnerTorch.R
mlr_learners_torch.Rd

This base class provides the basic functionality for training and prediction of a neural network. All torch learners should inherit from this class.

Validation

To specify the validation data, you can set the $validate field of the Learner, which can be set to:

  • NULL: no validation

  • ratio: only proportion 1 - ratio of the task is used for training and ratio is used for validation.

  • "test" means that the "test" task of a resampling is used and is not possible when calling $train() manually.

  • "predefined": This will use the predefined $internal_valid_task of a mlr3::Task.

This validation data can also be used for early stopping, see the description of the Learner's parameters.

Saving a Learner

In order to save a LearnerTorch for later usage, it is necessary to call the $marshal() method on the Learner before writing it to disk, as the object will otherwise not be saved correctly. After loading a marshaled LearnerTorch into R again, you then need to call $unmarshal() to transform it into a useable state.

Early Stopping and Internal Tuning

In order to prevent overfitting, the LearnerTorch class supports early stopping via the patience and min_delta parameters, see the Learner's parameters. When tuning a LearnerTorch it is also possible to combine the explicit tuning via mlr3tuning and the LearnerTorch's internal tuning of the epochs via early stopping. To do so, you just need to include epochs = to_tune(upper = <upper>, internal = TRUE) in the search space, where <upper> is the maximally allowed number of epochs, and configure the early stopping.

Checkpointing and Resuming

It is possible to save intermediate results from a run via the t_clbk("checkpoint") callback. It is then possible to train for more epochs by setting the resume parameter of the LearnerTorch. This parameter can either be a path or TRUE which will use the path of the provided checkpoint callback. Only the number of epochs should be changed between resumed runs, other parameter changes are considered undefined behavior. Also, make sure to use the same train-validation split. When the latest written checkpoint was for n1 epochs, the learner needs to be configured to be trained for n >= n1 epochs and the training will run for n2 = n - n1 epochs. With n = n1 the checkpointed run is already finished, so nothing is trained and the model of the checkpoint is returned – which is what lets a script that restarts itself be run again after it succeeded, and what recovers the model of a run that was killed after its last epoch. Configuring n < n1 is an error. Resuming will load the network weights, optimizer states and callback states. For some callbacks, training for n1 and then n2 epochs via resuming is not the same as training for n epochs from the start. This is for example the case for learning rate schedulers that depend on the total number of epochs to train for. The callbacks document their behavior under a corresponding Resuming section in their documentation. Furthermore, rng states are not restored, which constitutes another difference between a full and a resumed training run.

Network Head and Target Encoding

Torch learners are expected to have the following output:

  • binary classification: (batch_size, 1), representing the logits for the positive class.

  • multiclass classification: (batch_size, n_classes), representing the logits for all classes.

  • regression: (batch_size, 1) representing the response prediction.

A network may return more than one tensor, in which case it returns a list() of them. There are two typical reasons for this:

  • Networks with auxiliary classifiers such as Inception v3 return additional predictions that only exist to contribute to the loss during training. Here, every element has the shape given above and the first one is the prediction of interest.

  • A prediction that consists of several quantities – e.g. a mean and a standard deviation – is expressed by returning one tensor per quantity, which the prediction encoding combines.

In both cases the complete output is what the rest of the learner works with:

  • The loss is applied to it. Because the configured loss expects a single tensor, a learner whose network returns a list has to wrap it by overloading .loss_fn(), see the list of methods below. ContextTorch makes the output available as ctx$y_hats.

  • The prediction is encoded from it, both when predicting and when calculating the training and validation scores, so .encode_prediction() always receives the complete network output. Such a learner needs an encode_prediction() method for its task type, or has to overload the private .encode_prediction() method, because the encodings of the built-in task types expect a single tensor.

Note that the complete output is whatever the network returned in the mode it was called in, so a network whose extra tensors exist only during training – as auxiliary classifiers do – returns a different structure during training than during prediction, and .encode_prediction() has to handle both. classif.inception_v3 does this by encoding only the prediction of the main classifier.

Furthermore, the target encoding is expected to be as follows:

  • regression: The numeric target variable of a TaskRegr is encoded as a torch_float with shape c(batch_size, 1).

  • binary classification: The factor target variable of a TaskClassif is encoded as a torch_float with shape (batch_size, 1) where the positive class (Task$positive, which is also ensured to be the first factor level) is 1 and the negative class is 0.

  • multi-class classification: The factor target variable of a TaskClassif is a label-encoded torch_long with shape (batch_size) where the label-encoding goes from 1 to n_classes.

Predicting Tensors

The predict type "lazy_tensor", available for the task type "torch", hands back what the network produced – a lazy_tensor with one element per observation – instead of asking the task's default_encoder to turn it into a response. It is how to get at the logits of a classifier or the reconstruction of an autoencoder, and a task predicted this way needs no encoder at all. Unlike "prob" and "se" it is not opt-in: every learner for this task type has it among its $predict_types, because handing the output back does not depend on how the learner encodes a prediction. Set learner$predict_type = "lazy_tensor" to use it. A network with more than one head hands back one lazy_tensor per head, held in a data.table with one column per head so that the prediction is still one row per observation; as.data.table() spreads it into lazy_tensor.<head> columns.

Two things to know before using it:

  • Such a prediction does not survive saveRDS(). It holds torch tensors, which are external pointers: saving succeeds, and the object then fails with external pointer is not valid the next time the tensors are touched, in this session or in another. This applies to a ResampleResult holding one as well – its row ids and scores survive, its tensors do not.

  • Nothing about it is lazy. A lazy_tensor built from a tensor holds that tensor, so a prediction of this type is the network's output in memory, and resample() holds every fold's – combining the folds concatenates them, since lazy tensors from different networks share no data descriptor and cannot be concatenated lazily.

Important Runtime Considerations

There are a few hyperparameters settings that can have a considerable impact on the runtime of the learner. These include:

  • device: Use a GPU if possible.

  • num_threads: Set this to the number of CPU cores available if training on CPU. When resampling, benchmarking or tuning in parallel, each worker uses num_threads threads, so divide the available cores among the workers instead to avoid oversubscribing the machine.

  • tensor_dataset: Set this to TRUE (or "device" if on a GPU) if the dataset fits into memory. This loads and stacks every batch once up front, so it must not be used with a lazy_tensor that applies random data augmentation – the augmentation would then be drawn only once.

  • batch_size: Especially for very small models, choose a larger batch size.

  • batch_size_predict: Prediction has no backward pass and so fits larger batches than training.

  • jit_trace: Set this to TRUE to remove the per-batch R interpreter overhead, but only for a network whose control flow and shapes do not depend on the data. See the parameter's own entry.

  • num_workers: Load batches in parallel worker processes.

  • pin_memory: With a GPU, this speeds up the host-to-device copy of each batch.

Also, see the Early Stopping and Internal Tuning section for how to terminate training early.

Model

The Model is a list of class "learner_torch_model" with the following elements:

  • network :: The trained network.

  • optimizer :: The $state_dict() optimizer used to train the network.

  • loss_fn :: The $state_dict() of the loss used to train the network.

  • callbacks :: The callbacks used to train the network.

  • seed :: The seed that was / is used for training and prediction.

  • epochs :: How many epochs the model was trained for (early stopping).

  • task_col_info :: A data.table() containing information about the train-task.

Parameters

General:

The parameters of the optimizer, loss and callbacks, prefixed with "opt.", "loss." and "cb.<callback id>." respectively, as well as:

  • epochs :: integer(1)
    The number of epochs.

  • device :: character(1)
    The device. One of "auto", "cpu", or "cuda" or other values defined in mlr_reflections$torch$devices. The value is initialized to "auto", which will select "cuda" if possible, then try "mps" and otherwise fall back to "cpu".

  • num_threads :: integer(1)
    The number of threads for intraop parallelization (if device is "cpu"). This value is initialized to 1. When resampling, benchmarking or tuning in parallel, each worker uses this many threads, so divide the available cores among the workers instead of setting this to the number of cores.

  • num_interop_threads :: integer(1)
    The number of threads for interop parallelization (if device is "cpu"). Note that this can only be set once per session, so setting this for one learner also changes the behavior of other learners, and a later learner asking for a different value errors. NULL (default) uses whatever is set. In order to use different values for this parameter, use encapsulation to train the learners in separate R sessions.

  • seed :: integer(1) or "random" or NULL
    The torch seed that is used during training and prediction. This value is initialized to "random", which means that a random seed will be sampled at the beginning of the training phase. This seed (either set or randomly sampled) is available via $model$seed after training and used during prediction. Note that by setting the seed during the training phase this will mean that by default (i.e. when seed is "random"), clones of the learner will use a different seed. If set to NULL, no seeding will be done. This parameter only seeds torch's random number generator, it does not seed R's. Anything that is drawn from R's RNG is therefore unaffected by it, so to make those parts reproducible you need to seed R's RNG as well, e.g. via set.seed().

  • tensor_dataset :: logical(1) | "device"
    Whether to load all batches at once at the beginning of training and stack them. This is initialized to FALSE. If set to "device", the device of the tensors will be set to the value of device, which can avoid unnecessary moving of tensors between devices. When your dataset fits into memory this will make the loading of batches faster. Note that this should not be set for datasets that contain lazy_tensors with random data augmentation, as this augmentation will only be applied once at the beginning of training.

  • jit_trace :: logical(1)
    Whether to trace the network with torch::jit_trace() once at the start of training and then train the traced module instead of the original one. Not all learners support this.

Evaluation:

  • measures_train :: Measure or list() of Measures
    Measures to be evaluated during training.

  • measures_valid :: Measure or list() of Measures
    Measures to be evaluated during validation.

  • eval_freq :: integer(1)
    How often the train / validation predictions are evaluated using measures_train / measures_valid. This is initialized to 1. Note that the final model is always evaluated.

Resuming:

  • resume :: character(1) or TRUE
    Continues training from a checkpoint written by t_clbk("checkpoint"), either the folder it wrote to or TRUE, which takes that folder from the checkpoint callback of this learner. Note that epochs is the total number of epochs, i.e. it includes the epochs the checkpoint was already trained for: resuming a checkpoint from epoch 5 with epochs = 8 trains 3 more epochs.

Early Stopping:

  • patience :: integer(1)
    This activates early stopping using the validation scores. If the performance of a model does not improve for patience evaluation steps, training is ended. Note that this counts evaluation steps, not epochs: when eval_freq is greater than 1, patience evaluation steps correspond to patience * eval_freq epochs. Note that the final model is stored in the learner, not the best model, unless restore_best_weights is set to TRUE. This is initialized to 0, which means no early stopping. The first entry from measures_valid is used as the metric. This also requires to specify the $validate field of the Learner, as well as measures_valid. If this is set, the epoch after which no improvement was observed, can be accessed via the $internal_tuned_values field of the learner, and the validation scores of that epoch via its $best_valid_scores field.

  • min_delta :: double(1)
    The minimum improvement threshold for early stopping. Is initialized to 0.

  • restore_best_weights :: logical(1)
    Whether to restore the weights of the best epoch when training ends, instead of keeping those of the last epoch that was trained. Like min_delta, this only has an effect when early stopping is active, i.e. when patience is greater than 0. Is initialized to FALSE, i.e. the network of the last epoch is stored. Setting this to TRUE makes the stored network the one of the epoch that $internal_tuned_values reports, and costs one additional copy of the network's parameters in memory. Because $internal_valid_scores describes the network that is stored, it then reports the scores of the best epoch, i.e. the same scores as $best_valid_scores – except on a resumed run that never beats the score its checkpoint had already reached, which remembers no weights to restore and so still ends on those of its last epoch. Checkpoints written by t_clbk("checkpoint") are unaffected: they always hold the network as training left it.

Dataloader:

  • batch_size :: integer(1)
    The batch size used by the training and prediction dataloader. It is required for training (unless a batch_sampler is provided, which already determines the batches) and it is required for prediction (unless batch_size_predict is set).

  • batch_size_predict :: integer(1)
    The batch size used by the prediction dataloader (this includes the validation data during training). When set, it overrides batch_size for prediction. The batch size does not change the predictions, but smaller batches take longer and require less memory.

  • shuffle :: logical(1)
    Whether to shuffle the instances in the dataset. This is initialized to TRUE, which differs from the default (FALSE). It is ignored when a sampler or batch_sampler is provided.

  • sampler :: torch::sampler
    Object that defines how the dataloader draws samples, i.e. the order in which the observations are drawn. This must be the sampler generator (as returned by torch::sampler()), not an instance, as it is instantiated with the training dataset internally.

  • batch_sampler :: torch::sampler
    Object that defines how the dataloader draws batches. As for sampler, this must be the generator. When it is provided, the parameters batch_size, shuffle and drop_last are ignored during training, because the batch sampler already determines the batches.

  • num_workers :: integer(1)
    The number of workers for data loading (batches are loaded in parallel). The default is 0, which means that data will be loaded in the main process.

  • collate_fn :: function
    How to merge a list of samples to form a batch.

  • pin_memory :: logical(1)
    Whether the dataloader copies tensors into CUDA pinned memory before returning them.

  • drop_last :: logical(1)
    Whether to drop the last training batch in each epoch during training. Default is FALSE. It is ignored when a batch_sampler is provided.

  • timeout :: numeric(1)
    The timeout value for collecting a batch from workers. Negative values mean no timeout and the default is -1.

  • worker_init_fn :: function(id)
    A function that receives the worker id (in [1, num_workers]) and is executed after seeding on the worker but before data loading.

  • worker_globals :: list() | character()
    When loading data in parallel, this makes it possible to export globals to the workers. If this is a character vector, the objects in the global environment with those names are copied to the workers.

  • worker_packages :: character()
    Which packages to load on the workers.

Also see torch::dataloader for more information.

Inheriting

There are no separate classes for classification and regression to inherit from. Instead, the task_type must be specified as a construction argument. Any task type that is registered in mlr_reflections$task_types can be used. Support for a task type that mlr3torch does not know is added by implementing methods for the three S3 generics that hold the task-type-specific behaviour: output_dim_for() (how many output neurons the network needs), get_target_batchgetter() (how the target is turned into a tensor) and encode_prediction() (how the network's output is turned back into a prediction). Such a learner also has to be given a loss explicitly. This class can also be used for custom task types, see TaskTorch and the Custom Learning Problems article for more information.

When inheriting from this class, one should overload the following methods:

  • .network(task, param_vals)
    (Task, list()) -> nn_module
    Construct a torch::nn_module object for the given task and parameter values, i.e. the neural network that is trained by the learner. Note that a specific output shape is expected from the returned network, see section Network Head and Target Encoding. That section also describes when a network can return more than one tensor. You can use output_dim_for() to obtain the correct output dimension for a given task.

  • .loss_fn(task, param_vals)
    (Task, list()) -> nn_module
    Construct the loss that is applied to the output of the network. The default implementation generates the loss that was configured by the user, i.e. self$loss$generate(task). Overload this if the network returns more than one prediction and the configured loss has to be wrapped, see the aux_logits parameter of classif.inception_v3.

  • .ingress_tokens(task, param_vals)
    (Task, list()) -> named list() with TorchIngressTokens
    Create the TorchIngressTokens that are passed to the task_dataset constructor. The number of ingress tokens must correspond to the number of input parameters of the network. If there is more than one input, the names must correspond to the inputs of the network. See ingress_num, ingress_categ, and ingress_ltnsr on how to easily create the correct tokens. For more flexibility, you can also directly implement the .dataset(task, param_vals) method, see below.

  • .dataset(task, param_vals)
    (Task, list()) -> torch::dataset
    Create the dataset for the task. Don't implement this if the .ingress_tokens() method is defined. The dataset must return a named list where:

    • x is a list of torch tensors that are the input to the network. For networks with more than one input, the names must correspond to the inputs of the network.

    • y is the target tensor.

    • .index are the indices of the batch (integer() or a torch_int()).

    For information on the expected target encoding of y, see section Network Head and Target Encoding. Moreover, one needs to pay attention respect the row ids of the provided task. It is recommended to relu on task_dataset for creating the dataset.

It is also possible to overwrite the private .dataloader() method. This must respect the dataloader parameters from the ParamSet.

  • .dataloader(dataset, param_vals)
    (dataset, list()) -> torch::dataloader
    Create a dataloader from the dataset. Needs to respect at least batch_size and shuffle (otherwise predictions will be incorrectly ordered). Use get_batch_size(param_vals, "train") to obtain the batch size for the respective phase, which takes the batch_size_predict parameter into account.

To change the predict types, it is possible to overwrite the method below:

  • .encode_prediction(network_output, task)
    (torch_tensor or list() of them, Task) -> list()
    Take in the raw predictions from self$network (network_output) and encode them into a format that can be converted to valid mlr3 predictions using mlr3::as_prediction_data(). It is a list() of tensors when the network returns more than one, see section Network Head and Target Encoding. This method must take self$predict_type into account.

While it is possible to add parameters by specifying the param_set construction argument, it is currently not possible to remove existing parameters, i.e. those listed in section Parameters. None of the parameters provided in param_set can have an id that starts with "loss.", "opt.", or "cb.", as these are preserved for the dynamically constructed parameters of the optimizer, the loss function, and the callbacks.

To perform additional input checks on the task, the private .check_train_task(task, param_vals) and .check_predict_task(task, param_vals) can be overwritten. These should return TRUE if the input task is valid and otherwise a string with an error message.

For learners that have other construction arguments that should change the hash of a learner, it is required to implement the private $.additional_phash_input().

See also

Other Learner: GraphLearnerTorch, as_learner_torch(), mlr_learners.ft_transformer, mlr_learners.mlp, mlr_learners.module, mlr_learners.tab_resnet, mlr_learners.tabm, mlr_learners.torch_featureless, mlr_learners_torch_image, mlr_learners_torch_model

Super class

mlr3::Learner -> LearnerTorch

Active bindings

validate

How to construct the internal validation data. This parameter can be either NULL, a ratio in $(0, 1)$, "test", or "predefined".

loss

(TorchLoss)
The torch loss.

optimizer

(TorchOptimizer)
The torch optimizer.

callbacks

(list() of TorchCallbacks)
List of torch callbacks. The ids will be set as the names.

internal_valid_scores

Retrieves the internal validation scores of the epoch that the stored network comes from, as a named list(). This is the last epoch, unless restore_best_weights is TRUE and the best weights were actually restored, in which case it is the best epoch and these scores are the same as $best_valid_scores. A resumed run whose epochs never beat the score its checkpoint had already reached remembers no weights to restore, so the two fields describe different epochs even then. Specify the $validate field and the measures_valid parameter to configure this. Returns NULL if learner is not trained yet.

best_valid_scores

Retrieves the internal validation scores of the best epoch as a named list(). This is the epoch that is also reported via $internal_tuned_values, i.e. the epoch with the best score of the first validation measure. Unless restore_best_weights is TRUE, the trained network is the one after the last epoch, so this can differ from $internal_valid_scores. Tracking the best epoch requires early stopping to be active (patience > 0), so this is an empty list when it is not, as well as when the learner was trained without validation data – no early stopping callback runs in either case, and the two are not told apart. Returns NULL if the learner is not trained yet, and also when the model was not stored: unlike $internal_valid_scores, this is not part of the Learner contract that mlr3 snapshots into $state, so it can only be read back off the model. After a resample() or benchmark() with store_models = FALSE it is therefore NULL even though $internal_tuned_values still names the epoch it would describe.

internal_tuned_values

When early stopping is active, this returns a named list with the early-stopped epochs, otherwise an empty list is returned. Returns NULL if learner is not trained yet.

marshaled

(logical(1))
Whether the learner is marshaled.

network

(nn_module())
Shortcut for learner$model$network.

param_set

(ParamSet)
The parameter set

hash

(character(1))
Hash (unique identifier) for this object.

phash

(character(1))
Hash (unique identifier) for this partial object, excluding some components which are varied systematically during tuning (parameter values).

Methods

Public methods

Inherited methods

LearnerTorch$new()

Creates a new instance of this R6 class.

Usage

LearnerTorch$new(
  id,
  task_type,
  param_set,
  properties = character(),
  man,
  label,
  feature_types,
  optimizer = NULL,
  loss = NULL,
  packages = character(),
  predict_types = NULL,
  callbacks = list(),
  jittable = FALSE
)

Arguments

id

(character(1))
The id for of the new object.

task_type

(character(1))
The task type.

param_set

(ParamSet or alist())
Either a parameter set, or an alist() containing different values of self, e.g. alist(private$.param_set1, private$.param_set2), from which a ParamSet collection should be created.

properties

(character())
The properties of the object. See mlr_reflections$learner_properties for available values.

man

(character(1))
String in the format [pkg]::[topic] pointing to a manual page for this object. The referenced help package can be opened via method $help().

label

(character(1))
Label for the new instance.

feature_types

(character())
The feature types. See mlr_reflections$task_feature_types for available values, Additionally, "lazy_tensor" is supported.

optimizer

(NULL or TorchOptimizer)
The optimizer to use for training. Defaults to adam.

loss

(NULL or TorchLoss)
The loss to use for training. Defaults to MSE for regression and cross entropy for classification. For other task types there is no default and the loss has to be given, because which loss is appropriate depends on the learning problem.

packages

(character())
The R packages this object depends on.

predict_types

(character())
The predict types. See mlr_reflections$learner_predict_types for available values. For regression, the default is "response". For classification, this defaults to "response" and "prob". For the task type "torch", it defaults to "response". For other task types, it defaults to all predict types that are registered for the task type. To deviate from the defaults, it is necessary to overwrite the private $.encode_prediction() method, see section Inheriting.

callbacks

(list() of TorchCallbacks)
The callbacks to use for training. Defaults to an empty list(), i.e. no callbacks. Within a stage they are called in the order in which they are provided, unless a callback requests otherwise via its $weight, see section Ordering of CallbackSet.

jittable

(logical(1))
Whether the model can be jit-traced. Default is FALSE.

LearnerTorch$format()

Helper for print outputs.

Usage

LearnerTorch$format(...)

Arguments

...

(ignored).

LearnerTorch$print()

Prints the object.

Usage

LearnerTorch$print(...)

Arguments

...

(any)
Currently unused.

LearnerTorch$marshal()

Marshal the learner.

Usage

LearnerTorch$marshal(...)

Arguments

...

(any)
Additional parameters.

Returns

self

LearnerTorch$unmarshal()

Unmarshal the learner.

Usage

LearnerTorch$unmarshal(...)

Arguments

...

(any)
Additional parameters.

Returns

self

LearnerTorch$dataset()

Create the dataset for a task.

Usage

LearnerTorch$dataset(task)

Arguments

task

Task
The task

Returns

dataset

LearnerTorch$clone()

The objects of this class are cloneable with this method.

Usage

LearnerTorch$clone(deep = FALSE)

Arguments

deep

Whether to make a deep clone.