Retrieve the best accuracy from outputs of reference script

get_best_acc.py

#!/usr/bin/env python3

import argparse
import os
from typing import Tuple

def _extract(line:str) -> Tuple[float, float]:
    s = line.split(" ")
    acc = float(s[3])
    acc5 = float(s[5])
    return (acc, acc5)


def best_of_job(user: str, experiment_folder:str, job_id: int) -> Tuple[float, float, int, float, float, int]:
    out_file = f"{experiment_folder}/{user}/experiments/{job_id}/{job_id}_0_log.out"

    cur_epoch = None
    best_acc = None
    best_acc5 = None
    best_acc_epoch = None
    best_ema_acc = None
    best_ema_acc5 = None
    best_ema_acc_epoch = None

    with open(out_file, "r") as f:
        for line in f.readlines():
            if line.startswith("Epoch: ["):
                cur_epoch = int(line.split(" ")[1][1:-1])
            elif line.startswith("Test:  Acc@1"):
                acc, acc5 = _extract(line)
                if best_acc is None or best_acc < acc:
                    best_acc = acc
                    best_acc5 = acc5
                    best_acc_epoch = cur_epoch
            elif line.startswith("Test: EMA Acc@1"):
                s = line.split(" ")
                ema_acc, ema_acc5 = _extract(line)
                if best_ema_acc is None or best_ema_acc < ema_acc:
                    best_ema_acc = ema_acc
                    best_ema_acc5 = ema_acc5
                    best_ema_acc_epoch = cur_epoch
    return (best_acc, best_acc5, best_acc_epoch, best_ema_acc, best_ema_acc5, best_ema_acc_epoch)

def main(user: str, experiment_folder: str, job_ids: str) -> None:
    job_ids = [int(x) for x in job_ids.split(",")]
    best_acc = None
    best_acc5 = None
    best_acc_epoch = None
    best_ema_acc = None
    best_ema_acc5 = None
    best_ema_acc_epoch = None

    for job_id in job_ids:
        acc, acc5, epoch, ema_acc, ema_acc5, ema_acc_epoch = best_of_job(user, experiment_folder, job_id)
        if best_acc is None or best_acc < acc:
            best_acc = acc
            best_acc5 = acc5
            best_acc_epoch = epoch
        if best_ema_acc is None or best_ema_acc < ema_acc:
            best_ema_acc = ema_acc
            best_ema_acc5 = ema_acc5
            best_ema_acc_epoch = ema_acc_epoch
    print("Best acc@1\tacc@5\tepoch")
    print(f"{best_acc}\t{best_acc5}\t{best_acc_epoch}")
    print("Best EMA acc@1\tacc@5\tepoch")
    print(f"{best_ema_acc}\t{best_ema_acc5}\t{best_ema_acc_epoch}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--job-ids", type=str, required=True, help="job ids seperated by comma.")
    parser.add_argument("--user", type=str, help="user name.")
    parser.add_argument("--experiment-folder", type=str, default="/data/checkpoints", help="experiment folder that contains output files.")
    args = parser.parse_args()
    if args.user is None:
        args.user = os.environ["USER"]

    main(args.user, args.experiment_folder, args.job_ids)