Module `pretraining`

This IGM module trains a neural network emulator for the iceflow solver offline, before running a full glacier simulation. Given a precomputed dataset of input–velocity pairs, it optimises the emulator weights using a combination of a data loss (matching reference velocities) and a physics loss (based on the ice-flow energy functional), then writes the trained model to disk ready to be loaded by the iceflow module.

New emulator generation

This module trains the new generation of IGM emulators introduced in IGM 3.2, which use a different architecture and training pipeline from the emulators shipped with earlier versions. Older pre-trained emulators remain fully supported for inference (forward runs) and can still be loaded, just ensure that cfg.processes.iceflow.unified.network.pretrained_path is empty and IGM will fallback to trying to load a previous-generation emulator. However, they cannot be retrained or fine-tuned with this updated module.

Note

the pretraining module is intended for training new ice-flow emulators, if you want to use these new-generation emulators for inference, just set cfg.processes.iceflow.unified.network.pretrained_path to point to an existing emulator file and cfg.processes.iceflow.unified.network.pretrained to true.

Overview

The emulator learns to map input fields (e.g. ice thickness, surface elevation, sliding coefficient) to the horizontal velocity fields \(U\) and \(V\). Training adjusts network weights to minimize the sum of two loss functions:

Data loss — penalises the difference between predicted and reference velocities from a physics-based solver. Either mean-squared error (MSE) or a Huber loss can be used; Huber is recommended for glacier datasets which may contain large velocity outliers near the margins.
Physics loss — penalises departure from the ice-flow energy functional. This acts as a regulariser and can improve generalisation to conditions not present in the training data.

The relative weight of the physics loss, \(\lambda\), is adapted automatically during training: after an initial warmup phase (data-only), the trainer compares the gradient magnitudes of both losses and adjusts \(\lambda\) so that neither term dominates. An EMA filter and per-step multiplicative clamp prevent large oscillations.

Module ordering

The pretraining module must be listed alongside iceflow in the processes list. The iceflow module provides the network architecture, numerical discretisation (\(N_z\), precision), and physics configuration that the trainer uses. When pretraining is present, iceflow skips its normal initialisation of glacier fields.

defaults:
  - override /assimilations:
    - pretraining
  - override /processes:
    - iceflow

Training data

Training data must be provided as TFRecord files organised in the following directory layout:

<data_dir>/
  metadata.json
  train/
    nz<Nz>/
      shard_000.tfrecord
      shard_001.tfrecord
      ...
  val/
    nz<Nz>/
      shard_000.tfrecord
      ...

where <Nz> matches cfg.processes.iceflow.numerics.Nz. The number and order of input channels encoded in the TFRecords must match cfg.processes.iceflow.unified.inputs exactly. The module validates this at startup and raises a descriptive error if there is a mismatch.

Public dataset

A glacier catalogue in the required TFRecord format will be made publicly available following the publication of the paper describing this training pipeline.

Output

After training, the module writes a self-contained emulator artifact to:

<out_dir>/<experiment_name>/

This directory contains the Keras model file and input normaliser state. The iceflow module can load it directly by pointing cfg.processes.iceflow.unified.network.pretrained_path at this path.

When save_model: true (the default), training checkpoints are also written to <out_dir>/<experiment_name>/checkpoints/ so that interrupted runs can be resumed (see resume).

When make_plots: true, a loss-curve figure (loss_curve.png) and per-epoch speed-comparison images are written to <out_dir>/<experiment_name>/figures/.

Parameters

Default configuration file (pretraining.yaml):

pretraining:
  data_dir: "/path/to/tfrecords/"
  out_dir: "/path/to/save/models/"
  experiment_name: "name_of_model"

  # --- Batching / accumulation ---
  batch_size: 8 # the batch size is the number of training images you want to process before updating the model weights.
  micro_batch_size: 8 # here a micro batch is the amount of training images you want to process simultaneously on the GPU. If your GPU memory is not large enough to process the full batch size, you can set this to a smaller number and the training loop will automatically accumulate gradients until it has processed the full batch size before updating the model weights. Note that the batch size must be divisible by the micro batch size.

  # --- Training schedule ---
  epochs: 1000
  steps_per_epoch: 1000  # optimizer updates per epoch
  val_steps: 50          # val mini-batches evaluated per epoch
  learning_rate: 0.0001
  split_seed: 0          # seed used to shuffle the list of train shards

  # --- Loss ---
  loss_type: "huber"
  huber_delta: 50.0      # only used when loss_type == "huber"

  # --- Adaptive physics weighting (lambda) ---
  warmup_steps: 100000        # physics loss disabled for this many optimizer updates
  lambda_update_every: 200    # recompute lambda every N optimizer steps post-warmup
  lambda_ema: 0.99            # EMA factor when adopting the new lambda
  lambda_min: 0.001           # lower clip on lambda
  lambda_max: 100.0           # upper clip on lambda
  lambda_max_change: 1.2      # symmetric per-update multiplicative clamp; growth per epoch ≈ ×1.01 (×1.002^5 updates), ~380 epochs to scale 2000×

  # --- I/O ---
  resume: false
  make_plots: false
  save_model: true # if false do not save the model -OR- any checkpoints (useful for hyperparameter tuning)

Description of the parameters:

Name	Description	Default value	Units
`data_dir`	Path to the TFRecord dataset root directory. Must contain metadata.json and train/ and val/ subdirectories organised by Nz.	/path/to/tfrecords/	—
`out_dir`	Root directory under which the experiment output folder is created.	/path/to/save/models/	—
`experiment_name`	Name of the subdirectory created inside out_dir that holds the trained model, checkpoints, and figures.	name_of_model	—
`batch_size`	Effective batch size: the number of training examples processed before each optimizer weight update. Must be divisible by micro_batch_size.	8	—
`micro_batch_size`	Number of examples processed simultaneously on the GPU per forward-backward pass. Gradients are accumulated over batch_size / micro_batch_size micro-steps before each update. Reduce this if GPU memory is insufficient for the full batch size.	8	—
`epochs`	Total number of training epochs.	1000	—
`steps_per_epoch`	Number of optimizer weight updates performed per epoch.	1000	—
`val_steps`	Number of validation mini-batches evaluated at the end of each epoch.	50	—
`learning_rate`	Learning rate for the Adam optimizer.	0.0001	—
`split_seed`	Integer seed used to shuffle the list of training shards at the start of training. Set to a fixed value for reproducibility.	0	—
`loss_type`	Data loss function applied to predicted vs. reference velocities. Options: huber (recommended; robust to outliers near glacier margins) or mse.	huber	—
`huber_delta`	Transition point of the Huber loss: residuals smaller than delta are penalised quadratically, larger residuals linearly. Only used when loss_type is huber.	50.0	m y\( ^{-1} \)
`warmup_steps`	Number of optimizer updates at the start of training during which the physics loss is disabled. The network learns from data alone during this phase, which stabilises early training.	100000	—
`lambda_update_every`	After the warmup phase, the physics loss weight lambda is recomputed every this many optimizer steps by comparing data and physics gradient norms.	200	—
`lambda_ema`	Exponential moving average factor applied when updating lambda. Values close to 1 produce slow, stable adaptation; lower values respond faster but with more variance.	0.99	—
`lambda_min`	Lower bound on the adaptive physics loss weight lambda.	0.001	—
`lambda_max`	Upper bound on the adaptive physics loss weight lambda.	100.0	—
`lambda_max_change`	Multiplicative clamp applied to each lambda update: the new value is restricted to [lambda / c, lambda * c] where c is this parameter. Prevents large sudden changes.	1.2	—
`resume`	Resume training from the latest checkpoint in out_dir/experiment_name/checkpoints/. The experiment directory must already exist. Has no effect when save_model is false.	False	—
`make_plots`	Write a loss curve figure and per-epoch speed-comparison plots to out_dir/experiment_name/figures/ during training.	False	—
`save_model`	Save the trained emulator artifact and training checkpoints to disk. Set to false during hyperparameter searches to avoid writing many models.	True	—

Example usage

Basic training run

defaults:
  - override /processes:
    - iceflow
    - pretraining

processes:
  iceflow:
    method: unified
    numerics:
      Nz: 10
      precision: single
    unified:
      inputs: [thk, usurf, tau_ref]
      mapping: network
      network:
        pretrained: false
        architecture: CNN

  pretraining:
    data_dir: /data/glacier_catalog/tfrecords/
    out_dir: /experiments/
    experiment_name: cnn_nz10_v1
    batch_size: 8
    micro_batch_size: 8
    epochs: 1000
    steps_per_epoch: 1000
    val_steps: 50
    learning_rate: 0.0001
    loss_type: huber
    huber_delta: 50.0
    warmup_steps: 100000
    save_model: true
    make_plots: true

Resuming an interrupted training run

Set resume: true to pick up from the last saved checkpoint. All other parameters (epochs, learning rate, etc.) should remain unchanged.

  pretraining:
    data_dir: /data/glacier_catalog/tfrecords/
    out_dir: /experiments/
    experiment_name: cnn_nz10_v1
    # ... same hyperparameters as original run ...
    resume: true
    save_model: true

Hyperparameter search (no model saving)

When running many trials, set save_model: false to skip writing checkpoints and model artifacts. A state.score value (mean validation loss over the last 5 epochs) is still available for the search framework.

  pretraining:
    data_dir: /data/glacier_catalog/tfrecords/
    out_dir: /experiments/sweeps/
    experiment_name: trial_42
    batch_size: 16
    micro_batch_size: 4
    epochs: 50
    steps_per_epoch: 500
    val_steps: 20
    learning_rate: 0.0003
    loss_type: mse
    warmup_steps: 10000
    save_model: false
    make_plots: false

Memory-constrained GPU (gradient accumulation)

If a full batch does not fit in GPU memory, reduce micro_batch_size. The effective gradient update is equivalent to training with batch_size examples; only the peak GPU memory footprint changes.

  pretraining:
    batch_size: 16
    micro_batch_size: 2

Contributors: Sebastian Rosier.

Module pretraining