Function save_top_N_models

fmpose3d/common/utils.py:300–376

Save a checkpoint if it belongs to the top-N best (by lower data_threshold). Maintains an index file ' _top_models.json' in save_dir. Returns the path of the last saved checkpoint if a new one was saved, otherwise returns previous_name unchanged.

(
    previous_name,
    save_dir,
    epoch,
    data_threshold,
    model,
    model_name,
    num_saved_models=3,
)

def save_top_N_models(
    previous_name,
    save_dir,
    epoch,
    data_threshold,
    model,
    model_name,
    num_saved_models=3,
):
    """
    Save a checkpoint if it belongs to the top-N best (by lower data_threshold).
    Maintains an index file '<model_name>_top_models.json' in save_dir.

    Returns the path of the last saved checkpoint if a new one was saved,
    otherwise returns previous_name unchanged.
    """
    os.makedirs(save_dir, exist_ok=True)
    ckpt_path = os.path.join(
        save_dir, f"{model_name}_{epoch}_{int(data_threshold * 100)}.pth"
    )
    index_path = os.path.join(save_dir, f"{model_name}_top_models.json")

    # load current list
    top_list = []
    if os.path.exists(index_path):
        with open(index_path, "r") as f:
            top_list = json.load(f)

    # decide if we should save
    should_save = False
    if len(top_list) < int(num_saved_models):
        should_save = True
    else:
        worst_item = max(top_list, key=lambda x: x.get("p1", float("inf")))
        if data_threshold < float(worst_item.get("p1", float("inf"))):
            should_save = True

    if not should_save:
        return previous_name

    # save new checkpoint
    torch.save(model.state_dict(), ckpt_path)

    # append and trim to N
    top_list.append(
        {"p1": float(data_threshold), "path": ckpt_path, "epoch": int(epoch)}
    )
    # sort ascending by p1 and keep best N
    top_list.sort(key=lambda x: x.get("p1", float("inf")))
    while len(top_list) > int(num_saved_models):
        removed = top_list.pop()  # remove worst (last after sort ascending)
        if os.path.exists(removed.get("path", "")):
            os.remove(removed["path"])

    # write back index
    with open(index_path, "w") as f:
        json.dump(top_list, f, indent=2)