test.py

import torch
import torch.backends.cudnn as cudnn
import logging
import os

import numpy as np
from datetime import datetime
from cfgs.default_conf import cfg, load_cfg_fom_args
from common.utils import progress_bar, set_seed

from datasets.utils import build_dataloader
from models.utils import build_model
from algorithms.utils import build_algorithm


# get config file and logger
logger = logging.getLogger(__name__)
description = 'TRAI lab code demo'
load_cfg_fom_args(description)

# set seed for training
seed = set_seed(cfg.TRAIN.SEED)

# build dataset and dataloader
train_loader, test_loader = build_dataloader(cfg=cfg)

# load model and send to device
model = build_model(cfg=cfg)
model = model.cuda()

# check GPU ID and use them
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = torch.nn.DataParallel(model)
    cudnn.benchmark = True

# # load pretrained checkpoint or not. If you need add module in state_dict.
# if cfg.MODEL.CKPT_DIR is not None:
#     print('==> Resuming from checkpoint..')
#     checkpoint = torch.load(cfg.MODEL.CKPT_DIR)
#     # Add "module." prefix to state_dict keys
#     new_state_dict = {}
#     for k, v in checkpoint['state_dict'].items():
#         name = 'module.' + k 
#         new_state_dict[name] = v

#     # Load modified state_dict
#     model.load_state_dict(new_state_dict)
# else:
#     raise ValueError("You need a pretrain model")



# load pretrained checkpoint or not
if cfg.MODEL.CKPT_DIR is not None:
    print('==> Resuming from checkpoint..')
    checkpoint = torch.load(cfg.MODEL.CKPT_DIR)

    # Load modified state_dict
    model.load_state_dict(checkpoint['state_dict'])
else:
    raise ValueError("You need a pretrain model")

# build trainer
trainer = build_algorithm(
    cfg = cfg,
    model = model,
)

# test model
test_loss, acc = trainer.test(test_loader)
logger.info(f"Eval! loss: {test_loss.data.item():.4f}, accuracy rate: {acc:.2f}")