How to save checkpoints to Weights and baises ?

[sandeepgadhwal]

How to save checkpoints to Weights and baises ?

[zhouzaida]

Hi, you can implement a UploadCheckpointHook that is responsible for uploading checkpoints to WandB.

from mmengine.hooks import Hook
from mmengine.registry import HOOKS

@HOOKS.register_module()
class UploadCheckpointHook(Hook):
    def after_train_epoch(runner):
        wandb = runner.visualizer.get_backend('WandbVisBackend').experiment
        # https://docs.wandb.ai/ref/python/save#docusaurus_skipToContent_fallback
        wandb.save('checkpoint.pth', policy='now')

The below is snippets code:

runner = Runner(
    model=MMResNet50(),
    work_dir='./work_dir',
    train_dataloader=train_dataloader,
    optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)),
    train_cfg=dict(by_epoch=True, max_epochs=5, val_interval=1),
    val_dataloader=val_dataloader,
    val_cfg=dict(),
    val_evaluator=dict(type=Accuracy),
    # 1. set WandbVisBackend
    visualizer=dict(type='Visualizer', vis_backends=[dict(type='WandbVisBackend')]),
    # 2. set UploadCheckpointHook
    custom_hooks = [dict(type='UploadCheckpointHook')]
)

runner.train()

[pha123661]

Hi,

Thanks for your response, which is really useful!
I wonder if it is possible to get the name of the best checkpoint in the custom hook? since its name is dynamic.
ex: "${work_dir}/best_acc_top1_epoch_14.pth"

Thank you.

[saarthakk-insitro]

Following code works too if you want to save in artifacts in WandB:

from mmengine.hooks import Hook
from mmengine.registry import HOOKS

@HOOKS.register_module()

class WandbCheckpointHook(Hook):

def __init__(self, interval=16000, artifact_name="model"):
    self.interval = interval
    self.artifact_name = artifact_name
    self.run_id = None
    self.version = 0

def before_run(self, runner):
    # Get the existing wandb run ID
    self.run_id = wandb.run.id if wandb.run else None
    if self.run_id is None:
        raise RuntimeError("Wandb run not initialized. Please call wandb.init() before using this hook.")

def after_train_iter(self, runner, batch_idx, data_batch=None, outputs=None):
    if self.every_n_iters(runner, self.interval):
        self.upload_checkpoint(runner)

def upload_checkpoint(self, runner):
    filename = f'iter_{runner.iter + 1}.pth'
    print(filename)
    filepath = os.path.join(runner.work_dir, filename)
    print(filepath)
    runner.save_checkpoint(runner.work_dir, filename=filename)
    
    # Create a new artifact
    artifact = wandb.Artifact(
        name=f"{self.artifact_name}-{self.run_id}",
        type="model",
        description=f"Model checkpoint for iter {runner.iter + 1}"
    )
    
    # Add file to artifact
    artifact.add_file(filepath, name=filename)
    
    # Log artifact to wandb with specific version
    wandb.log_artifact(artifact, aliases=[f"v{self.version}"])
    
    print(f"Checkpoint {filename} uploaded to wandb as artifact {self.artifact_name}-{self.run_id}:v{self.version}")
    
    # Increment version for next checkpoint
    self.version += 1

def every_n_iters(self, runner, n):
    return (runner.iter + 1) % n == 0 if n > 0 else False

And in config add following:

custom_hooks = [dict(type='WandbCheckpointHook')]
cfg.custom_hooks = custom_hooks