checkpointing.py

import os
import typing
from typing import Any

import torch


def save_checkpoint(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    iteration: int,
    out: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
    extra_state: dict[str, Any] | None = None,
):
    state_dict = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "iteration": iteration,
    }
    if extra_state is not None:
        state_dict["extra_state"] = extra_state
    torch.save(state_dict, out)


def load_checkpoint(
    src: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer | None,
    extra_state: dict[str, Any] | None = None,
):
    state_dict = torch.load(src, map_location="cpu")
    model.load_state_dict(state_dict["model"])
    if optimizer is not None:
        optimizer.load_state_dict(state_dict["optimizer"])
        _move_optimizer_state_to_device(optimizer)
    if extra_state is not None:
        extra_state.clear()
        extra_state.update(state_dict.get("extra_state", {}))
    return state_dict["iteration"]


def _move_optimizer_state_to_device(optimizer: torch.optim.Optimizer) -> None:
    for param_group in optimizer.param_groups:
        for param in param_group["params"]:
            state = optimizer.state.get(param)
            if not state:
                continue
            for key, value in list(state.items()):
                if torch.is_tensor(value):
                    state[key] = value.to(device=param.device)