모델 학습 후 성능 최적화를 위해 하이퍼파라미터 튜닝을 수행할 때가 있다. 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_mAPrun_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()으로 전달한다.
'Artificial Intelligence > 머신러닝-딥러닝' 카테고리의 다른 글
| [Pytorch 오류해결] RuntimeError: stack expects each tensor to be equal size, but got (0) | 2022.11.16 |
|---|---|
| 파이토치 분산 학습하기 (Pytorch Distributed Training) (0) | 2022.11.12 |
| 파이토치 모델 서빙하고 내보내기 (도커, 토치서브, 토치스크립트, ONNX 활용) (0) | 2022.11.10 |
| [자연어처리] 챗봇 시스템 구축 (Rasa 활용) (0) | 2022.11.08 |
| [신경망 이론] 신경망 학습과 손실함수 (Neural Network Training And Loss Function) (0) | 2022.03.17 |
