smrunner.train Namespace Reference

Functions
	adjust_learning_rate (optimizer, epoch, lr, adjust_epoch, adjust_rate)
	tmp_pipeline (sample, keys, device)
	train_sample (model, loss_cls, optimizer, sample)
	train_epoch (model, loss_cls, optimizer, data_loader, collect_keys, device)
	validation_sample (model, loss_cls, sample, metric_args)
	validation (model, loss_cls, data_loader, collect_keys, device, metric_args)
	train (cfg, device_ids)

Function Documentation

adjust_learning_rate()

smrunner.train.adjust_learning_rate	(		optimizer,
			epoch,
			lr,
			adjust_epoch,
			adjust_rate )

Definition at line 51 of file train.py.

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
 
#데이터 전처리, 학습 디바이스로 데이터 전송

tmp_pipeline()

smrunner.train.tmp_pipeline	(		sample,
			keys,
			device )

Definition at line 64 of file train.py.

def tmp_pipeline(sample, keys, device):
    for key in keys:
        sample[key] = sample[key].to(device)
 
    return sample
 
#배치에 대해 학습

train()

smrunner.train.train	(		cfg,
			device_ids )

Definition at line 142 of file train.py.

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()} )")
 

train_epoch()

smrunner.train.train_epoch	(		model,
			loss_cls,
			optimizer,
			data_loader,
			collect_keys,
			device )

Definition at line 85 of file train.py.

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
 
#배치에 대해 평가

train_sample()

smrunner.train.train_sample	(		model,
			loss_cls,
			optimizer,
			sample )

Definition at line 71 of file train.py.

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 학습 기능

validation()

smrunner.train.validation	(		model,
			loss_cls,
			data_loader,
			collect_keys,
			device,
			metric_args )

Definition at line 119 of file train.py.

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를 이용해서 학습함

validation_sample()

smrunner.train.validation_sample	(		model,
			loss_cls,
			sample,
			metric_args )

Definition at line 102 of file train.py.

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()
 
#전체 평가 데이터에 대한 평가