kaggle-ckpt 0.1.0

Production checkpoint management for Kaggle and ephemeral training environments

kaggle-ckpt

Production checkpoint management for Kaggle and other ephemeral training environments.

Kaggle notebooks wipe /kaggle/temp between sessions and cap /kaggle/working at roughly 20 GB. kaggle-ckpt handles the lifecycle for you: fast atomic saves during training, configurable retention so you don't blow the disk quota, and a single call at the end to promote your best model (plus metrics, logs, and config) into persistent storage.

Features

• Automatic /kaggle/temp (fast, ephemeral) → /kaggle/working (persistent) promotion
• Atomic saves via PID-scoped temp files, so a crash mid-write never corrupts a checkpoint
• Configurable retention: keep the best, the last, and/or the last N periodic checkpoints
• Exponential Moving Average (EMA) of model weights, with serialization built in
• Robust state-dict loading that auto-handles DataParallel's module. prefix
• Kaggle / Colab / local environment auto-detection with local fallback paths
• Framework-agnostic: works with raw PyTorch, Lightning, Hugging Face, etc.

Install

pip install kaggle-ckpt

Quickstart

from kaggle_ckpt import CheckpointManager, CheckpointConfig

cfg = CheckpointConfig("my_experiment", monitor="val_macro_f1", mode="max")
ckpt = CheckpointManager(cfg)

for epoch in range(1, num_epochs + 1):
    train_metrics = train_one_epoch(...)
    val_metrics = validate(...)

    state = {
        "epoch": epoch,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
    }
    ckpt.save_epoch(epoch, state, metrics=val_metrics)

# At the end of training, promote everything to /kaggle/working
ckpt.copy_final_artifacts(
    metrics=test_metrics,
    train_log=ckpt.metrics.get_history(),
    class_names=class_names,
    config_dict=config,
)

Loading a checkpoint later

from kaggle_ckpt import CheckpointManager, CheckpointConfig

ckpt = CheckpointManager(CheckpointConfig("my_experiment"))
checkpoint = ckpt.load_best()
ckpt.load_model_weights(model, checkpoint=checkpoint, use_ema=True)

EMA

from kaggle_ckpt import EMA

ema = EMA(model, decay=0.9998)

for step in training_loop:
    optimizer.step()
    ema.update()

ema.apply_shadow()
validate(model)
ema.restore()

License

MIT