aiaccel.torch.lightning.callbacks.LoadPretrainedCallback

class aiaccel.torch.lightning.callbacks.LoadPretrainedCallback(model_path: Path, target_patterns: list[str], pattern_map: dict[str, str] | None = None, source_excludes: list[str] | None = None, target_excludes: list[str] | None = None, config_name: str = 'merged_config.yaml')

Initialize a model from a pretrained checkpoint before training or validation.

The callback loads weights from model_path once fitting or validation begins, matches finetune parameters to pretrained ones using glob-like patterns, and copies the matching weights into the finetune module before any optimization steps run.

Parameters:

• model_path – Directory containing checkpoints saved by load_checkpoint().

• target_patterns – Glob expressions that describe finetune parameters which should be initialized from pretrained weights.

• pattern_map – Optional mapping from finetune patterns to pretrained ones. Wildcards ("*") are allowed and must appear the same number of times on both sides of the mapping.

• source_excludes – Optional pretrained-side glob patterns that should never be copied even when referenced by pattern_map.

• target_excludes – Optional finetune-side glob patterns that should never be overwritten.

• config_name – Name of the checkpoint configuration to load.

Example:

callback = LoadPretrainedCallback(
    model_path=Path("pretrain_ckpt"),
    target_patterns=["detr_module.*"],
    pattern_map={"backbone.*": "visual_backbone.*"},
    source_excludes=["detr_module.heads.cls_head.*"],
    config_name="merged_config.yaml",
)
trainer = lt.Trainer(callbacks=[callback])
trainer.fit(model)

__init__(model_path: Path, target_patterns: list[str], pattern_map: dict[str, str] | None = None, source_excludes: list[str] | None = None, target_excludes: list[str] | None = None, config_name: str = 'merged_config.yaml') → None

Methods

__init__(model_path, target_patterns[, ...])
load_state_dict(state_dict)	Called when loading a checkpoint, implement to reload callback state given callback's state_dict.
on_after_backward(trainer, pl_module)	Called after loss.backward() and before optimizers are stepped.
on_before_backward(trainer, pl_module, loss)	Called before loss.backward().
on_before_optimizer_step(trainer, pl_module, ...)	Called before optimizer.step().
on_before_zero_grad(trainer, pl_module, ...)	Called before optimizer.zero_grad().
on_exception(trainer, pl_module, exception)	Called when any trainer execution is interrupted by an exception.
on_fit_end(trainer, pl_module)	Called when fit ends.
on_fit_start(trainer, pl_module)	Load pretrained weights and copy them into matching finetune parameters.
on_load_checkpoint(trainer, pl_module, ...)	Called when loading a model checkpoint, use to reload state.
on_predict_batch_end(trainer, pl_module, ...)	Called when the predict batch ends.
on_predict_batch_start(trainer, pl_module, ...)	Called when the predict batch begins.
on_predict_end(trainer, pl_module)	Called when predict ends.
on_predict_epoch_end(trainer, pl_module)	Called when the predict epoch ends.
on_predict_epoch_start(trainer, pl_module)	Called when the predict epoch begins.
on_predict_start(trainer, pl_module)	Called when the predict begins.
on_sanity_check_end(trainer, pl_module)	Called when the validation sanity check ends.
on_sanity_check_start(trainer, pl_module)	Called when the validation sanity check starts.
on_save_checkpoint(trainer, pl_module, ...)	Called when saving a checkpoint to give you a chance to store anything else you might want to save.
on_test_batch_end(trainer, pl_module, ...[, ...])	Called when the test batch ends.
on_test_batch_start(trainer, pl_module, ...)	Called when the test batch begins.
on_test_end(trainer, pl_module)	Called when the test ends.
on_test_epoch_end(trainer, pl_module)	Called when the test epoch ends.
on_test_epoch_start(trainer, pl_module)	Called when the test epoch begins.
on_test_start(trainer, pl_module)	Called when the test begins.
on_train_batch_end(trainer, pl_module, ...)	Called when the train batch ends.
on_train_batch_start(trainer, pl_module, ...)	Called when the train batch begins.
on_train_end(trainer, pl_module)	Called when the train ends.
on_train_epoch_end(trainer, pl_module)	Called when the train epoch ends.
on_train_epoch_start(trainer, pl_module)	Called when the train epoch begins.
on_train_start(trainer, pl_module)	Called when the train begins.
on_validation_batch_end(trainer, pl_module, ...)	Called when the validation batch ends.
on_validation_batch_start(trainer, ...[, ...])	Called when the validation batch begins.
on_validation_end(trainer, pl_module)	Called when the validation loop ends.
on_validation_epoch_end(trainer, pl_module)	Called when the val epoch ends.
on_validation_epoch_start(trainer, pl_module)	Called when the val epoch begins.
on_validation_start(trainer, pl_module)	Ensure pretrained weights are loaded before running validation.
setup(trainer, pl_module, stage)	Called when fit, validate, test, predict, or tune begins.
state_dict()	Called when saving a checkpoint, implement to generate callback's state_dict.
teardown(trainer, pl_module, stage)	Called when fit, validate, test, predict, or tune ends.

Attributes

state_key

Identifier for the state of the callback.