Function load_checkpoint

utils/common.py:1590–1678

恢复训练的checkpoint加载函数 Args: checkpoint_path (str): checkpoint文件路径 model (torch.nn.Module): 主模型 ema_model (torch.nn.Module): EMA模型 optimizer (torch.optim.Optimizer, optional): 优化器 lr_scheduler (torch.optim.lr_scheduler._LRScheduler, optional): 学

(
        checkpoint_path,
        ema_model,
        optimizer=None,
        lr_scheduler=None,
        strict=True,
        print_info=True,
        map_location="cuda"
)

    return total_params

def load_checkpoint(
        checkpoint_path,
        ema_model,
        optimizer=None,
        lr_scheduler=None,
        strict=True,
        print_info=True,
        map_location="cuda"
):
    """
    恢复训练的checkpoint加载函数

    Args:
        checkpoint_path (str): checkpoint文件路径
        model (torch.nn.Module): 主模型
        ema_model (torch.nn.Module): EMA模型
        optimizer (torch.optim.Optimizer, optional): 优化器
        lr_scheduler (torch.optim.lr_scheduler._LRScheduler, optional): 学习率调度器
        map_location (str, optional): 加载设备

    Returns:
        cfg (dict): 保存的配置字典
        global_step (int): 当前训练步数
    """
    print(f"Loading checkpoint from {checkpoint_path} ...")
    checkpoint = torch.load(checkpoint_path, map_location=map_location)

    # 恢复配置
    cfg = checkpoint["cfg"]

    # 加载模型参数
    weights = checkpoint["weights"]
    ema_weights = checkpoint["ema_weights"]
    load_info_online = ema_model.online_model.load_state_dict(weights, strict=strict)
    load_info_ema = ema_model.ema_model.load_state_dict(ema_weights, strict=strict)
    print(f"✅ EMA weights loaded : online_weights {weights_num(weights)}, ema_weights : {weights_num(ema_weights)}")

    print_load_report(load_info_online, "online_model", len(weights), print_info=print_info)
    print_load_report(load_info_ema, "ema_model", len(ema_weights), print_info=print_info)

    # 加载优化器和学习率调度器
    if optimizer is not None and "optimizer" in checkpoint:

        for i, param_group in enumerate(optimizer.param_groups):
            print(f"optimizer参数load之前：Param group {i}:")
            for key, value in param_group.items():
                if key != "params":  # params 太长
                    print(f"  {key}: {value}")

        optimizer.load_state_dict(checkpoint["optimizer"])

        for i, param_group in enumerate(optimizer.param_groups):
            print(f"optimizer参数load之后：Param group {i}:")
            for key, value in param_group.items():
                if key != "params":  # params 太长
                    print(f"  {key}: {value}")

        removed_cnt = 0