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.

MetricSummaryCallback

Bases: Callback

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

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 train 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.

PlottingCallback

Bases: Callback

A callback to create plots 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

uncertainty_variables = {'ice_conc': 'total_standard_uncertainty'} instance-attribute

plotter = Plotter(DEFAULT_SIC_SPEC + plot_spec) instance-attribute

plotter_metadata = None 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_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 SIC uncertainty in the same normalised scale as the target.

make_plots(trainer, pl_module, dataset, batch_size)

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.

set_metadata(config, model_name)

Set metadata for the plotter.

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.