pytorch: optuna ddp 적용하는 방법

DDP optuna 적용 방법

DDP에서 적용하는 방법을 알아볼게요.
생각보다 간단하답니다.

먼저, DDP 학습을 위해서 초기화를 해주어야 해요.

dist.init_process_group("gloo")
rank = dist.get_rank()

그런 다음 study.optimize를 하나의 gpu만 실행할 수 있도록 if rank == 0 구문으로 나누어줘요.

study = None
n_trials = 20
if rank == 0:
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=n_trials)
else:
    for _ in range(n_trials):
        try:
            objective(None)
        except optuna.TrialPruned:
            pass

마지막으로 objective에서 trial = optuna.integration.TorchDistributedTrial(single_trial)을 해주면 끝이에요.

def objective(single_trial):
    trial = optuna.integration.TorchDistributedTrial(single_trial)

    # Generate the model.
    model = DDP(define_model(trial).to(DEVICE))
    
    ...

전체 코드를 보고 싶다면 참고에 있는 링크를 클릭해주세요!

참고

https://github.com/optuna/optuna-examples/blob/main/pytorch/pytorch_distributed_simple.py