Skip to contents

Checkpoint Callback

Source: R/CallbackSetCheckpoint.R
mlr_callback_set.checkpoint.Rd

Saves the optimizer, weights, and callback states every freq epochs as well as the final state. This can be used to later continue a training run via the resume parameter of LearnerTorch.

Checkpoints are written at the end of an epoch. For one written after epoch <n>, three files are created in path:

  • network<n>.pt :: The $state_dict() of the network.

  • optimizer<n>.pt :: The $state_dict() of the optimizer.

  • state<n>.rds :: The epoch, the version of mlr3torch that wrote the checkpoint, the $state_dict()s of the training run's other callbacks, so that a later run can continue, as well as some other information. Additionally, there is run.rds which contains some additioanl global meta information.

Details

Saving the learner itself in the callback with a trained model is impossible, as the model slot is set after the last callback step is executed.

Resuming

This callback is special because it enables resuming a training run. Its own state is the folder it writes to, which learner$model$callbacks$<id>$path reports – the only way to learn where a path function sent a run, e.g. one fit of a resample(). That state is not part of a checkpoint and is not restored: a resuming run writes where its own path says.

Ordering

This callback has weight Inf and therefore runs last, so it captures all the changes other callbacks made.

See also

Other Callback: TorchCallback, as_torch_callback(), as_torch_callbacks(), callback_set(), mlr3torch_callbacks, mlr_callback_set, mlr_callback_set.progress, mlr_callback_set.tb, mlr_callback_set.unfreeze, mlr_context_torch, t_clbk(), torch_callback()

Super class

CallbackSet -> CallbackSetCheckpoint

Public fields

weight

(numeric(1))
Always Inf, see section Ordering.

Methods

Public methods

Inherited methods

CallbackSetCheckpoint$new()

Creates a new instance of this R6 class.

Usage

CallbackSetCheckpoint$new(path, freq)

Arguments

path

(character(1) | function())
The path to a folder where the models are saved, or a function of no arguments returning it. The latter is especially useful to create unique directories during resample() or benchmark() per fit. The folder must be new, empty, or already contain checkpoints. A half-written checkpoint – what a run killed mid-write leaves behind – may be written over, since a resuming run continues from the newest complete one.

freq

(integer(1))
How often the model is saved, in epochs.

CallbackSetCheckpoint$state_dict()

Returns the folder this callback writes to so it can be accessed from the learner when the path was a function.

Usage

CallbackSetCheckpoint$state_dict()

CallbackSetCheckpoint$on_begin()

Checks whether the checkpoint path is valid.

Usage

CallbackSetCheckpoint$on_begin()

CallbackSetCheckpoint$on_epoch_end()

Saves the network and optimizer state dict. Does nothing if freq is not met.

Usage

CallbackSetCheckpoint$on_epoch_end()

CallbackSetCheckpoint$on_end()

Saves the final network and optimizer, unless the last epoch was already saved.

Usage

CallbackSetCheckpoint$on_end()

CallbackSetCheckpoint$clone()

The objects of this class are cloneable with this method.

Usage

CallbackSetCheckpoint$clone(deep = FALSE)

Arguments

deep

Whether to make a deep clone.

Examples

cb = t_clbk("checkpoint", freq = 1)
task = tsk("iris")

pth = tempfile()
learner = lrn("classif.mlp", epochs = 3, batch_size = 1, callbacks = cb)
learner$param_set$set_values(cb.checkpoint.path = pth)

learner$train(task)

list.files(pth)
#>  [1] "network1.pt"   "network2.pt"   "network3.pt"   "optimizer1.pt"
#>  [5] "optimizer2.pt" "optimizer3.pt" "run.rds"       "state1.rds"   
#>  [9] "state2.rds"    "state3.rds"   

# continue training for 3 more epochs, starting from the last checkpoint
learner_resumed = lrn("classif.mlp", epochs = 6, batch_size = 1, resume = pth)
learner_resumed$train(task)
learner_resumed$model$epochs
#> [1] 6