ggml_trainer: A Training Context You Step Yourself

In ggmlR: 'GGML' Tensor Operations for Machine Learning

A Training Context You Step Yourself

Description

Builds the graph and optimizer context for a compiled model and hands them back, so a training loop can be written in R rather than handed to ggml_fit. The context lives until $free(), keeping weights, Adam moments and the compute graph across as many steps as the loop takes.

Usage

ggml_trainer(
  model,
  batch_size = 32L,
  nbatch_logical = NULL,
  loss_mask = NULL,
  weight_decay = NULL
)

Arguments

model	A compiled ggml_sequential_model, or a compiled functional model with one input and one output
batch_size	Samples per step. The graph is built for exactly this many; every batch handed to $step() must match.
nbatch_logical	Logical batch size for gradient accumulation. Defaults to batch_size, meaning every step updates the weights.
loss_mask	NULL (default) for the compiled loss unweighted; 1 for one weight per sample; the number of outputs for one weight per output element. See the section on loss weights.
weight_decay	Decoupled weight decay in [0, 1]. NULL (default) takes the value given to ggml_compile. Frozen weights are not decayed.

Details

This is the level below ggml_fit_opt_gen. Reach for it when the loop is not a pass over data: reinforcement learning, where experience comes from a replay buffer the agent is still filling and a target network is refreshed every so many steps; curriculum schemes that change the data distribution as they go; data generated fresh for every step; or anything where "epoch" is not the unit of progress.

Do not emulate this by calling ggml_fit once per batch: each such call builds a new optimizer context, so the Adam moments are reset and every step is the optimizer's first.

Value

An object of class ggml_trainer: a list of functions $step(x, y, sample_weight = NULL), $eval(x, y, sample_weight = NULL), $predict(x, batch_size = NULL, matmul_precision = NULL), $snapshot(), $update_target(target, tau = 1), $model(), $calibrate_bn(x), $set_lr(lr), $n_steps(), $free().

Using it

tr <- ggml_trainer(model, batch_size = 32L)
on.exit(tr$free())

for (i in seq_len(10000L)) {
  b <- sample_batch(replay, 32L)     # your own bookkeeping
  loss <- tr$step(b$x, b$y)
  if (i %% 500L == 0L) target <- tr$model()   # refresh a target network
}

model <- tr$model()                  # weights back into an R model object
tr$free()

What a step does

$step(x, y, sample_weight) runs one forward and backward pass and returns that step's loss. Gradient accumulation, if nbatch_logical asked for it, is tracked inside the context: the optimizer updates the weights only on the period boundary, and the intervening steps only accumulate. $eval(x, y, sample_weight) is the same without the backward pass, for measuring on held-out data.

Loss weights and masks

With loss_mask set, the loss is weighted mean squared error, sum(w * (pred - y)^2) / (batch * outputs), and $step() and $eval() take a sample_weight matrix:

• loss_mask = 1: sample_weight is [batch, 1] (or a vector of length batch) – one weight per sample, as in Keras.

• loss_mask = <outputs>: sample_weight is [batch, outputs] – one weight per output element. This is not Keras' sample_weight: a 0/1 matrix is a mask that trains only the outputs whose label is known (e.g. one action of a Q-network), and a 0 removes that element's gradient exactly. The loss is not renormalised by the number of active elements.

Omitting sample_weight on a weighted trainer means all ones. Only loss = "mse" supports weights.

Models

Sequential models and functional models (ggml_model) with one input and one output are supported. Functional models with several inputs or outputs are refused. For a functional model with dropout the dropout masks are redrawn before every $step().

Predicting during training

$predict(x, batch_size) returns predictions from the current weights, as ggml_predict would, in eval mode (no dropout). It runs on a scheduler of its own, created on first use: the weights are copied into a separate buffer for the call and released after it, so neither the training graph's buffers nor the optimizer state are touched. batch_size defaults to the trainer's; x may have any number of rows.

Target networks

$snapshot() returns an independent copy of the model: its own weight tensors in its own buffer and its own scheduler, sharing nothing mutable with the trainer. Later steps do not move it, ggml_predict on it never touches the training path, and it outlives $free(). It is built in eval mode, so dropout is off in it. Release it with ggml_snapshot_free (a finaliser is only a backstop); after that ggml_predict and $update_target on it are refused with an error.

$update_target(target, tau = 1) moves a snapshot towards the current weights, \theta' \leftarrow \tau\theta + (1 - \tau)\theta': tau = 1 is the hard copy of DQN, a small tau the Polyak average of DDPG/SAC. Non-trainable buffers (a sequential batch-norm's running mean and variance) are copied whole, never averaged. The target must have the same weights, by name and shape, or it is refused. The update goes through host memory, one read and one write per tensor.

Lifetime

The context holds GPU buffers and must be released with $free(); calling it twice is harmless. A trainer that goes out of scope without it is cleaned up when garbage collected, but that moment is not predictable, so free it yourself. After $free() every other method errors rather than touching released memory.

The model object passed in is unchanged – R semantics being what they are, the trainer works on its own copy of the weights. $model() returns a model carrying the current ones.

The returned model shares the trainer's backend scheduler, and predicting with it re-lays the scheduler's buffers for another graph. The trainer therefore re-allocates its own graph on the first step after any $model() call; weights and Adam moments are kept. $predict() needs none of this. Predicting with the original model object while the trainer is alive is not covered by this and can crash the session on the next step – take the model from $model() instead.

Batch norm

A step normalizes each batch by its own statistics and does not update the running ones a sequential batch_norm predicts with; ggml_fit() recalibrates them after training, a trainer does not. Call $calibrate_bn(x) with representative inputs (up to 1024 rows are used) before predicting, snapshotting or saving – or fix the statistics with ggml_freeze_bn_stats. Creating a trainer for a model with unfixed batch-norm layers warns once.

See Also

ggml_fit_opt_gen for a generator-driven epoch loop, ggml_fit for the ordinary in-memory path.

Other optimization: ggml_fit_opt(), ggml_fit_opt_gen(), ggml_fit_opt_multi(), ggml_opt_alloc(), ggml_opt_context_optimizer_type(), ggml_opt_dataset_data(), ggml_opt_dataset_free(), ggml_opt_dataset_get_batch(), ggml_opt_dataset_get_batch_head(), ggml_opt_dataset_init(), ggml_opt_dataset_labels(), ggml_opt_dataset_ndata(), ggml_opt_dataset_shuffle(), ggml_opt_dataset_weights(), ggml_opt_default_params(), ggml_opt_epoch(), ggml_opt_eval(), ggml_opt_fit(), ggml_opt_fit_multi(), ggml_opt_free(), ggml_opt_get_lr(), ggml_opt_grad_acc(), ggml_opt_init(), ggml_opt_init_for_fit(), ggml_opt_init_for_fit_multi(), ggml_opt_inputs(), ggml_opt_labels(), ggml_opt_loss(), ggml_opt_loss_i(), ggml_opt_loss_type_binary_cross_entropy(), ggml_opt_loss_type_cross_entropy(), ggml_opt_loss_type_huber(), ggml_opt_loss_type_mae(), ggml_opt_loss_type_mean(), ggml_opt_loss_type_mse(), ggml_opt_loss_type_sum(), ggml_opt_loss_type_weighted_mse(), ggml_opt_n_loss(), ggml_opt_ncorrect(), ggml_opt_optimizer_name(), ggml_opt_optimizer_type_adamw(), ggml_opt_optimizer_type_sgd(), ggml_opt_outputs(), ggml_opt_outputs_i(), ggml_opt_pred(), ggml_opt_prepare_alloc(), ggml_opt_reset(), ggml_opt_result_accuracy(), ggml_opt_result_free(), ggml_opt_result_init(), ggml_opt_result_loss(), ggml_opt_result_loss_i(), ggml_opt_result_n_loss(), ggml_opt_result_ndata(), ggml_opt_result_pred(), ggml_opt_result_reset(), ggml_opt_set_lr(), ggml_opt_static_graphs(), ggml_snapshot_free()

Examples

ggml_set_n_threads(1L)  # deterministic, single OpenMP pool
n <- 64
x <- matrix(runif(n * 4), nrow = n, ncol = 4)
y <- matrix(0, nrow = n, ncol = 2)
for (i in seq_len(n)) { y[i, if (sum(x[i, ]) > 2) 1L else 2L] <- 1 }

model <- ggml_model_sequential() |>
  ggml_layer_dense(8, activation = "relu") |>
  ggml_layer_dense(2, activation = "softmax")
model$input_shape <- 4L
model <- ggml_compile(model, optimizer = "adam",
                      loss = "categorical_crossentropy")

tr <- ggml_trainer(model, batch_size = 16L)
for (i in 1:4) {
  idx <- ((i - 1L) * 16L + 1L):(i * 16L)
  tr$step(x[idx, , drop = FALSE], y[idx, , drop = FALSE])
}
model <- tr$model()
tr$free()

ggmlR documentation built on Oct. 23, 2026, 1:07 a.m.