test.py

import torch

import argparse
import time, sys, os, yaml, json

from utils import evaluate
from models import load_model
from datasets import load_dataset
import wandb

def add_args_parser():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--ckpt_path', type=str)
    parser.add_argument('--use_wandb', action='store_true')
    
    return parser

def load_wandb(cfg):
    wandb.init(
        config=cfg,
        project='Torch Template (MNIST Classification)',
        group=f"test_{cfg['model']['name']}_{cfg['data']['test']['dataset']}"
    )

def main(args):
    ckpt = torch.load(args.ckpt_path, weights_only=False)
    cfg = ckpt['cfg']

    if args.use_wandb: 
        load_wandb(cfg)

    # Device Setting
    device = None
    if cfg['device'] != 'cpu' and torch.cuda.is_available():
        device = cfg['device']
    else: 
        device = 'cpu'
    print(f"device: {device}")
    
    # Hyperparameter Settings
    hp_cfg = cfg['hyperparameters']
    
    # Load Dataset
    data_cfg = cfg['data']['test']
    test_ds = load_dataset(data_cfg)
    test_dl = torch.utils.data.DataLoader(test_ds,
                                          batch_size=hp_cfg['batch_size'])
    print(f"Load Dataset {data_cfg['dataset']}")
    
    # Load Model
    model_cfg = cfg['model']
    model = load_model(model_cfg).to(device)
    model.load_state_dict(ckpt['model'])
    print(f"Load Model {model_cfg['name']}")
    print(f"It trained epochs: {ckpt['epoch']}")
    
    start_time = int(time.time())
    result_dict = evaluate(model, test_dl, device)
    test_time = int(time.time() - start_time)
    print(f"Test Time: {test_time//60:02d}m {test_time%60:02d}s")
    
    for key, value in result_dict.items():
        print(f"{key}: {value:.4f}")

    with open(os.path.join(os.path.dirname(args.ckpt_path), "result.json"), 'w') as f:
        json.dump(result_dict, f, indent=2)
    
    if args.use_wandb:
        wandb.log(result_dict)
        wandb.finish()

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Test', parents=[add_args_parser()])
    args, _ = parser.parse_known_args()

    main(args)