Skip to content

Callbacks

ActivationSaver

Bases: Callback

Register forward hooks and save captured activations to disk per batch.

The ActivationSaver plays two roles:

  1. A set of PyTorch forward hooks attached to named submodules of the model (encoders, processor.conv*, decoder stages, ...). Each hook stores a detached CPU copy of its module's output tensor for the first rollout step of each test batch.
  2. A Lightning Callback whose on_test_batch_end flushes the buffered activations (and, optionally, a copy of the raw batch tensors) to a single batch_{idx:05d}.pt file.

Parameters:

Name Type Description Default
layer_paths Sequence[str]

Dotted module paths resolvable via model.get_submodule (equivalently, keys of model.named_modules()).

 required
output_dir Path | str

Directory for per-batch .pt files and a metadata.json.

 required
save_inputs bool

If True, also save the raw batch tensors alongside the activations.

 True

BATCH_FILE_TEMPLATE = 'batch_{batch_idx:05d}.pt' class-attribute instance-attribute

METADATA_FILE = 'metadata.json' class-attribute instance-attribute

layer_paths = list(layer_paths) instance-attribute

output_dir = Path(output_dir) instance-attribute

save_inputs = save_inputs instance-attribute

attach(model)

Resolve layer paths on the model and register all forward hooks.

Raises:

Type Description
ValueError

If any requested layer path does not resolve to a submodule on the model.

detach()

Remove all registered forward hooks.

rollout_counter_increment(module, inputs)

Increment the rollout counter for this batch at each forward.

rollout_counter_reset(module, inputs)

Reset the processor rollout counter at the start of each batch.

on_test_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx=0)

Clear per-batch buffers and snapshot the raw batch tensors.

on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Persist captured activations for this batch to a single .pt file.

on_test_end(trainer, pl_module)

Write a metadata file and remove forward hooks when the test loop ends.

on_test_start(trainer, pl_module)

Called when the test begins.

EMAWeightAveragingCallback

Bases: WeightAveraging

A callback that updates an averaged model for Exponential Moving Average (EMA) after each training step.

every_n_epochs = every_n_epochs instance-attribute

every_n_steps = every_n_steps instance-attribute

on_train_batch_end(trainer, pl_module, *args, **kwargs)

Ignore the update if the module has no parameters.

on_train_epoch_end(trainer, pl_module)

Ignore the update if the module has no parameters.

should_update(step_idx=None, epoch_idx=None)

Update if we are at the requested number of steps or epochs.

MediaLoggingCallback

Bases: Callback

A callback to create and log images during evaluation.

frequency_batch = int((frequency or {}).get('batch', -1)) instance-attribute

frequency_epoch = int((frequency or {}).get('epoch', -1)) instance-attribute

frequency_number = int((frequency or {}).get('number', -1)) instance-attribute

make_input_plots = make_input_plots instance-attribute

make_static_plots = make_static_plots instance-attribute

make_video_plots = make_video_plots instance-attribute

prefix = prefix instance-attribute

cached_batch_idx_ = None instance-attribute

cached_dataloader_idx_ = None instance-attribute

cached_outputs_ = None instance-attribute

cache_batch(batch_idx, dataloader_idx, outputs)

Cache the current batch information for use in epoch-end plotting.

is_logged(kind, dates)

Record (kind, dataloader, date) as logged, returning True if already seen.

is_sample_batch(batch_idx, total_batches)

Return True if batch_idx is one of frequency_number equally-spaced targets.

load_dataset(dataloader)

Load the dataset for the given dataloader index.

load_target_uncertainties(dataset, dates)

Load uncertainty at the same normalised scale as the target.

make_plots(trainer, pl_module, dataset, batch_size)

on_batch_end(trainer, pl_module, outputs, batch_idx, dataloader_idx, *, is_last_batch, num_batches, dataloaders)

Shared test/validation batch-end handling: cache and maybe plot.

on_epoch_end(trainer, pl_module, dataloaders)

Shared test/validation epoch-end handling: maybe plot.

on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Called at the end of each test batch.

on_test_epoch_end(trainer, pl_module)

Called at the end of each test epoch.

on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Called at the end of each validation batch.

on_validation_epoch_end(trainer, pl_module)

Called at the end of each validation epoch.

MetricSummaryCallback

Bases: Callback

A callback to summarise metrics at the end of an epoch or a run.

STAGES_BY_TRAINER_FN = {TrainerFn.FITTING.value: ['train', 'validation'], TrainerFn.TESTING.value: ['test']} class-attribute

PLOT_GROUPS = {FractionalSkillScorePerForecastDay: 'fss', SpatialMeanGroundTruthPerForecastDay: 'spatial_mean', SpatialMeanPredictionPerForecastDay: 'spatial_mean'} class-attribute

log_per_epoch_metrics(trainer, metrics, stage)

Log per-epoch metrics to W&B.

log_per_run_metrics(trainer, metrics)

Log per-run metrics to W&B.

Note that these will be based on metrics accumulated during the final epoch, due to the reset behaviour in log_per_epoch_metrics.

on_test_epoch_start(trainer, pl_module)

Called at the start of a test epoch.

on_test_epoch_end(trainer, pl_module)

Called at the end of a test epoch.

on_train_epoch_start(trainer, pl_module)

Called at the start of a training epoch.

on_train_epoch_end(trainer, pl_module)

Called at the end of a training epoch.

on_validation_epoch_start(trainer, pl_module)

Called at the start of a validation epoch.

on_validation_epoch_end(trainer, pl_module)

Called at the end of a validation epoch.

teardown(trainer, pl_module, stage)

Called at the end of a run: log train/validation or test metrics.

PredictionWriter

Bases: Callback

Write evaluation predictions to a CF-style NetCDF file in the run directory.

enabled = enabled instance-attribute

output_path = None instance-attribute

mask_dir = None instance-attribute

on_test_start(trainer, pl_module)

Initialise prediction export before evaluation batches run.

on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)

Append one evaluation batch to the NetCDF file.

on_test_end(trainer, pl_module)

Close the NetCDF file after evaluation.

on_exception(trainer, pl_module, exception)

Close the NetCDF file if evaluation fails, leaving the partial output readable.

UnconditionalCheckpoint

Bases: Callback

A callback to summarise metrics during evaluation.

impl = ModelCheckpoint() instance-attribute

dirpath property writable

Return the directory path where checkpoints are saved.

on_train_end(trainer, pl_module)

Called when training ends.

save_unconditionally(trainer)

Save a checkpoint unconditionally.