test.py

Go to the documentation of this file.

 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
 
from smaction.datasets.action_dataset_loader import ActionDatasetLoader
from smutils.utils_os import create_directory, save_json
from smrunner.model_runner_builder import build_model
from smrunner.losses.loss_builder import build_loss
from smrunner.utils import *
from smrunner.metrics import *
import json
 
import datetime
import numpy as np
import copy
import os
 
__data_loader__ = {
    "ActionDatasetLoader": ActionDatasetLoader,
}
 
def tmp_pipeline(sample, keys, device):
    for key in keys:
        sample[key] = sample[key].to(device)
    # sample['keypoints'] = sample['keypoints'].to(device)
    # sample['label'] = sample['label'].to(device)
 
    return sample
 
 
def validation_sample(model, loss_cls, sample, device):
    
    tmp_keys = ['keypoints', 'label_action', 'label_pose']
    sample = tmp_pipeline(sample, tmp_keys, device)
    result = model(sample)
    sample.update(result)
    loss = loss_cls(sample)
    loss = torch.mean(loss)
 
    acc = top1_acc(sample, 'action_labels', 'label_action')
    return acc, loss.item()
 
 
def validation(model, loss_cls, data_loader, device):
    model.eval()
    max_step = len(data_loader)
    sum_loss = 0
    meter = AverageMeter()
    with torch.no_grad():
        for step, sample in enumerate(data_loader):
            val_acc, val_loss = validation_sample(model, loss_cls, sample, device)
            meter.update(val_acc)
            sum_loss += val_loss
            print(f'   val   => [{step+1}/{max_step}] -> top1 acc : {meter.mean()*100:.6f}%', end='\r')
        print('')
    return meter.mean(), sum_loss/max_step
 
 
def test(cfg, device_ids):
    num_workers = cfg.test.num_workers
    batch_size = cfg.test.batch_size
    
    if device_ids is None:
        device = 'cpu'
    else:
        device = f'cuda:{device_ids[0]}'
    
    # 데이터셋 로더 셋팅
    print('---- data loader setting ----')
    data_loader_cfg = copy.deepcopy(cfg.data_loader)
    DataSetClass = __data_loader__[data_loader_cfg.pop('type')]
    test_dataset = DataSetClass(mode='test', **data_loader_cfg)
    test_data_loader = DataLoader(test_dataset, batch_size, shuffle=False, num_workers=num_workers, drop_last=False)
    
    # 모델 빌드
    print('---- build model ----')
    model = build_model(cfg)
    model.to(device)
 
    load_model(cfg, model, cfg.test.pretrained , device=device)
    
    # 로스 빌드
    print('---- build loss ----')
    for key, loss_class in cfg.loss.items():
        loss_class.update(dict(device=device))
    loss_cls = build_loss(cfg)
    
    #테스트 결과 저장 폴더 생성
    create_directory(cfg.test.save_root)
 
    result_msg = []
    msg_path = os.path.join(cfg.test.save_root, 'test_result.json')
 
    print(f"\n-start- ( {datetime.datetime.now()} )")
 
    val_acc, val_loss = validation(model, loss_cls, test_data_loader, device)
 
    msg = f'val loss : {val_loss:.6f}   val acc : {val_acc*100:.6f}'
    print(msg)
 
    result_msg.append(msg)
    save_json(result_msg, msg_path)
 
 
        
    print(f"\n-end- ( {datetime.datetime.now()} )")