파이토치 Multi-GPU with Optuna

모델 학습 후 성능 최적화를 위해 하이퍼파라미터 튜닝을 수행할 때가 있다. Optuna 사용 방법은 여러 문서에 기술되어 있어 참고할 수 있지만 파이토치에서 분산처리를 적용한 Multi-GPU 환경에서 Optuna를 함께 실행하는 방법에 대한 레퍼런스가 적어 약간의 삽질을 수행했기 때문에 추후 참고 목적으로 작성한다. 또한 파이토치로 분산처리 적용하여 Multi-GPU 학습을 수행하고자할 때 mp.spawn의 리턴값을 일반적으로는 받아오지 못하는 경우가 있어 이에 대한 내용도 덧붙이고자 한다.

 

전체 의사코드는 아래와 같고, 각 함수마다 간단히 서술한다.

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup(rank, world_size):
    dist.init_process_group(backend='nccl', init_method="tcp://127.0.0.1:1234", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def run_tuning(tuning_fn, world_size, epochs, lr, optimizer_name):
    parent_conn, child_conn = mp.Pipe()
    mp.spawn(tuning_fn, args=(world_size, epochs, lr, optimizer_name, child_conn), nprocs=world_size, join=True)
    while parent_conn.poll():
        best_mAP = parent_conn.recv()
    return best_mAP
    
def tune_hyperparameter(rank, world_size, optimizer, scheduler, epochs, trial):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    model = Model().to(rank)
    model = DDP(model, device_ids=[rank])
    
    train_dataset = ''
    valid_dataset = ''
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
    valid_sampler = DistributedSampler(valid_dataset, num_replicas=world_size, rank=rank, shuffle=False)
    train_loader = torch.utils.data.DataLoader()
    valid_loader = torch.utils.data.DataLoader()

    for epoch in range(epochs):
        train()
        mAP = valid()
        ...

    conn.send(best_mAP)
    trial.report(mAP, epoch)
    if trial.should_prune():
        raise optuna.TrialPruned()

    cleanup()
    
def objective(trial, world_size):
    epochs = trial.suggest_int("epoch", 10, 20)
    lr = trial.suggest_float("lr", low=1e-7, high=1e-2, log=True)
    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "AdamW", "AdamP", "RAdam"])
    
    best_mAP = run_tuning(tune_hyperparameter, world_size, epochs, lr, optimizer_name)
    return best_mAP
    
    
if __name__ == "__main__":
	seed = 42
    ngpus_per_node = torch.cuda.device_count()
    world_size = ngpus_per_node

    func = lambda trial: objective(trial, world_size, world_size)

    study = optuna.create_study(direction="maximize", sampler=RandomSampler(seed=seed))
    study.optimize(func, n_trials=5)

    pruned_trials = study.get_trials(deepcopy=False, states=[TrialState.PRUNED])
    complete_trials = study.get_trials(deepcopy=False, states=[TrialState.COMPLETE])

    trial = study.best_trial.value

 

Optuna 사용 핵심은 objective 함수 작성이다. 함수 내에는 튜닝할 하이퍼파라미터의 종류와 그 값의 범위값을 작성해준다. 이후 학습 루틴을 수행 후 그 결과를 리턴값으로 전달해주는 것이 끝이다. 학습루틴을 objective 함수 내부에 정의해도 되나 분산학습을 함께 적용하기 위해 별도의 함수로 만들어 실행한다. 

import optuna

def objective(trial, world_size):
    epochs = trial.suggest_int("epoch", 10, 20)
    lr = trial.suggest_float("lr", low=1e-7, high=1e-2, log=True)
    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "AdamW", "AdamP", "RAdam"])
    
    best_mAP = run_tuning(tune_hyperparameter, world_size, epochs, lr, optimizer_name)
    return best_mAP
    
if __name__ == "__main__":
	func = lambda trial: objective(trial, world_size)
    study = optuna.create_study()
    study.optimize(func, n_trials=5)

run_tuning 함수는 인자로 실제로 하이퍼파라미터 튜닝을 수행할 함수를 받는다. run_tuning 함수는 다음과 같이 구성되어 있다. 

def run_tuning(tuning_fn, world_size, epochs, lr, optimizer_name):
    parent_conn, child_conn = mp.Pipe()
    mp.spawn(tuning_fn, args=(world_size, epochs, lr, optimizer_name, child_conn), nprocs=world_size, join=True)
    while parent_conn.poll():
        best_mAP = parent_conn.recv()
    return best_mAP

run_tuning은 실제로 하이퍼파라미터 튜닝을 수행할 함수를 첫 번째 인자(tuning_fn)로 받는다. mp.spawn은 리턴값이 없으므로 학습 관련 결과에 대한 값을 받아올 수 없다. 만약 multi-gpu를 통한 분산처리의 리턴값을 받아올 필요가 없다면 mp.spawn() 한 줄만 작성해주어도 된다. 하지만 필요한 경우가 있다. 이를 해결하기 위해 mp.Pipe()를 사용한다. 학습 관련 결과를 child_conn을 이용해 parent_conn으로 전송하는 것이다. child_conn은 학습 함수인 tuning_fn의 인자로 들어가서 학습 관련 결과에 대해 parent_conn으로 전달한다. mp.spawn이 종료되면 parent_conn에서 poll() 메서드와 recv() 메서드를 통해 그 결과를 가져올 수 있다. 

 

아래 함수는 실제로 파이토치에서 Multi-GPU를 활용해 학습하는 루틴을 정의한 함수다. 

def tune_hyperparameter(rank, world_size, optimizer, scheduler, epochs, trial):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    model = Model().to(rank)
    model = DDP(model, device_ids=[rank])
    
    train_dataset = ...
    valid_dataset = ...
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True)
    valid_sampler = DistributedSampler(valid_dataset, num_replicas=world_size, rank=rank, shuffle=False)
    train_loader = torch.utils.data.DataLoader(...)
    valid_loader = torch.utils.data.DataLoader(...)

    for epoch in range(epochs):
        train()
        mAP = valid()
        ...

    conn.send(max_best_mAP)
    trial.report(mAP, epoch)
    if trial.should_prune():
        raise optuna.TrialPruned()

    cleanup()

함수의 가장 처음과 끝은 setup()과 cleanup()으로 각각 Multi-GPU로 학습 가능하도록 초기화하고 학습 이후 환경 초기화를 수행하는 역할을 한다. 모델은 DDP로 래핑해주고 데이터셋도 DistributedSampler를 사용한 결과로 데이터로더를 만들어준다. 이후 학습을 수행하고 conn.send()를 통해 그 결과를 parent_conn()으로 전달한다.