train.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, ActionDatasetLoader_v22, ActionDatasetLoader_mtml
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
 
#TODO
#빌더로 빼서 따로 관리 필요할 것 같음 : 데이터 로더, 옵티마이저, 스케줄러
#데이터 로더를 어떻게 관리할지 결정 필요
#   현재 : 행동인식만 학습하기 때문에 smaction/dataset에 데이터로더를 구현해둠
#   변경 안 : 데이터 로더만 따로 모아서 관리
#       -> smrunner에 데이터로더 관련 기능들 정리
#       -> smdataset에 데이터로더 관련 기능들 정리
#학습 코드를 어떻게 관리할지 결정 필요
#   현재 : smrunner/train.py에 학습 코드가 구현되어 있음. 특별 케이스만 작동됨
#   변경 안: smrunner/trainer에 학습 관련 기능 구현(일반화시켜서)
#데이터 전처리 기능이 임시로 구현되어 있음 : tmp_pipeline
#   일반화 시켜서 따로 관리 필요함 : mmlab 참조해보기
#   현재 : 일부 전처리 기능은 데이터로더에서 작업됨, 디바이스로 전송하는 기능은 학습 코드에 임시로 작성됨
#학습 파라미터 저장 및 로드 구조 개선 필요
#멀티 gpu 학습이 가능하도록 코드 및 구조 개선 필요
 
 
#데이터 로더
__data_loader__ = {
    "ActionDatasetLoader" : ActionDatasetLoader,
    "ActionDatasetLoader_v22" : ActionDatasetLoader_v22,
    "ActionDatasetLoader_mtml" : ActionDatasetLoader_mtml,
}
 
#옵티마이저
__optimizer__ = {
    "SGD" : optim.SGD,
    "Adam" : optim.Adam
}
 
#lr 스케줄러(step)
def adjust_learning_rate(optimizer, epoch, lr, adjust_epoch, adjust_rate):
    #적용 lr 계산
    for ep, rate in zip(adjust_epoch, adjust_rate):
        if epoch >= ep:
            lr=rate*lr
    
    #lr 셋팅
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
        
    return lr
 
#데이터 전처리, 학습 디바이스로 데이터 전송
def tmp_pipeline(sample, keys, device):
    for key in keys:
        sample[key] = sample[key].to(device)
 
    return sample
 
#배치에 대해 학습
def train_sample(model, loss_cls, optimizer, sample):
    optimizer.zero_grad() #옵티마이저 기울기 초기화
 
    result = model(sample) #inference
    sample.update(result) #샘플에 inference 결과 추가
 
    loss = loss_cls(sample) #로스 계산
    loss = torch.mean(loss) #로스 평균
    loss.backward() #백프로파게이션, 기울기 계산
    optimizer.step() #모델 파라미터 업데이트
 
    return loss.item()
 
#1-epoch 학습 기능
def train_epoch(model, loss_cls, optimizer, data_loader, collect_keys, device):
    
    model.train() #모델 학습 모드로 변경
    max_step = len(data_loader) #몇번 업데이트하는지 확인, 출력용도
    sum_loss = 0 #로스 합, 출력용도
    for step, sample in enumerate(data_loader):
        
        sample = tmp_pipeline(sample, collect_keys, device) #데이터를 디바이스로 전송
        loss = train_sample(model, loss_cls, optimizer, sample) #샘플(배치) 훈련
 
        #진행도 화면 출력
        sum_loss += loss
        print(f"   train => [{step+1}/{max_step}] -> avg loss : {(sum_loss/(step+1)):.6f} ,  loss : {loss:.6f}", end='\r')
    print('')
    return sum_loss/max_step
 
#배치에 대해 평가
def validation_sample(model, loss_cls, sample, metric_args):
    
    result = model(sample) #inference
    sample.update(result) #샘플에 inference 결과 추가
    loss = loss_cls(sample) #로스 계산
    loss = torch.mean(loss) #로스 평균
 
    #top-1 정확도 계산
    # acc_dict = top1_acc_dict(sample, ['pred_action', 'pred_pose'], ['label_action', 'label_pose'])
    acc_dict = top1_acc_multi_task(sample, **metric_args)
                                #    pred_key=['pred_pose', 'pred_hand', 'pred_foot'], 
                                #    gt_key=['label', 'label', 'label'], 
                                #    target_tasks=['pose', 'hand', 'foot'],
                                #    task_key='category')
    return acc_dict, loss.item()
 
#전체 평가 데이터에 대한 평가
def validation(model, loss_cls, data_loader, collect_keys, device, metric_args):
    model.eval() #평가모드
    max_step = len(data_loader) #몇 스텝 반복하는지 확인, 출력용도
    sum_loss = 0 #로스 합, 출력용도
    meter = AverageMeterDict(model.predict_keys.keys()) #평균 계산기
 
    #평가
    with torch.no_grad():
        for step, sample in enumerate(data_loader):
 
            #평가용 데이터를 디바이스로 전송
            sample = tmp_pipeline(sample, collect_keys, device) 
 
            #배치에 대해 평가
            val_acc, val_loss = validation_sample(model, loss_cls, sample, metric_args)
            
            meter.update(val_acc) #평균 계산기에 정확도 계산 결과 업데이트
            sum_loss += val_loss #로스합
            print(f'   val   => [{step+1}/{max_step}] -> {meter.print_str()}', end='\r')
        print('')
    return meter, sum_loss/max_step
 
#학습 함수, 1-gpu를 이용해서 학습함
def train(cfg, device_ids):
    #변수 셋팅
    num_workers = cfg.train.num_workers
    init_lr = cfg.train.init_lr
    epochs = cfg.train.epochs
    batch_size = cfg.train.batch_size
    checkpoint_folder = os.path.join(cfg.train.save_root, 'weights')
 
    #학습 디바이스 셋팅
    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')] #사용할 데이터 로더 선택
    
    #학습용 데이터 로더 빌드
    train_dataset = DataSetClass(mode='train', **data_loader_cfg) 
    train_data_loader = DataLoader(train_dataset, batch_size, shuffle=True, num_workers=num_workers, drop_last=True)
    #평가용 데이터 로더 빌드
    test_dataset = DataSetClass(mode='val', **data_loader_cfg)
    test_data_loader = DataLoader(test_dataset, batch_size, shuffle=True, num_workers=num_workers, drop_last=True)
    
    # 모델 빌드
    print('---- build model ----')
    model = build_model(cfg)
    model.to(device)
    if cfg.train.pretrained is not None:
        load_model(cfg, model, cfg.train.pretrained , device=device) #사전학습파라미터 로드
    else:
        init_weights(model) #모델 파라미터 초기화
    
    # 로스 빌드
    print('---- build loss ----')
    for key, loss_class in cfg.loss.items():
        loss_class.update(dict(device=device))
    loss_cls = build_loss(cfg)
 
    # 옵티마이저 설정
    print('---- optimizer setting ----')
    assert cfg.train.optimizer in __optimizer__, \
            f'not found optimizer type : {cfg.train.optimizer}'
 
    optimizer = __optimizer__[cfg.train.optimizer](model.parameters(), lr=init_lr, **cfg.train.optimizer_args)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, **cfg.train.scheduler_args)
 
    #훈련 결과 저장 폴더 생성
    create_directory(checkpoint_folder)
 
    result_msg = []
    msg_path = os.path.join(cfg.train.save_root, 'train_result.json')
 
    #훈련 시작
    model.train()
    for ep in range(epochs):
        #lr 조정
        if cfg.train.scheduler == 'StepLR':
            lr = adjust_learning_rate(optimizer, ep, init_lr, cfg.train.adjust_lr_epoch, cfg.train.adjust_lr_rate)
        else:
            lr = scheduler.get_last_lr()
        print(f'[{ep+1}]/[{epochs}] -> time( {datetime.datetime.now()} ), lr({lr})')
        # print(f'[{ep+1}]/[{epochs}] -> time( {datetime.datetime.now()} )')
 
        #train
        avg_loss = train_epoch(model, loss_cls, optimizer, train_data_loader, cfg.collect_keys, device)
 
        #eval
        if (ep+1) % cfg.train.val_interval == 0:
            meter, val_loss = validation(model, loss_cls, test_data_loader, cfg.collect_keys, device, cfg.metric_args)
            
            if cfg.train.update_loss_weight and (ep+1) % cfg.train.update_loss_weight_interval == 0:
                val_acc = meter.mean()
                new_weights = [base_weight + (1-val_acc[key]) for key in pred_keys]
                loss_cls.weights = new_weights
 
            msg = f'[{ep+1}]/[{epochs}]    train loss : {avg_loss:.6f}   val loss : {val_loss:.6f}   {meter.print_str()}'
            print(msg)
 
            result_msg.append(msg)
            save_json(result_msg, msg_path)
 
        if cfg.train.scheduler != 'StepLR':
            scheduler.step()
 
        #모델 파라미터 저장
        checkpoint_path = os.path.join(checkpoint_folder, f'{ep}.pth')
        save_model(model, checkpoint_path)
 
    print(f"\n-end- ( {datetime.datetime.now()} )")