argparse, omegaconf로 실험 환경 구성하기

ML 프로젝트에서 다양한 실험을 진행하다 보면 다양한 하이퍼 파라미터들과 여러 설정값을 수정해야 하는 일이 자주 발생한다. 그럴 때마다 코드를 직접 수정하게 되면 수정해야 할 부분을 찾아야 하는 일도 번거롭고, 제대로 값을 수정하지 못하고 실험이 진행되는 일이 빈번하다. 따라서 이번 포스팅에서는 configuration을 적절하게 관리하는 방법에 대해서 알아보자.

---

argparse

argparse를 이용하면 CLI에서 argument의 값을 입력받을 수 있다. CLI에서 입력받은 argument 값을 사용하기 위해서는 코드를 다음과 같이 구성해야 한다.

1. parser 객체 생성
2. 입력받을 argument 선언
3. argument 파싱

아래는 예시 코드이다.

import argparse

def get_args():
    # parser 객체 생성
    parser = argparse.ArgumentParser(description='ML Experiment Configuration')
    
    # 입력받을 argument 선언
    parser.add_argument('--data_dir', type=str, required=True, help='Path to dataset directory')
    parser.add_argument('--model_type', type=str, required=True, choices=['cnn', 'lstm', 'bert'], help='Type of model to train')
    
    parser.add_argument('--lr', type=float, default=0.0001, help='Learning rate')
    parser.add_argument('--batch_size', type=int, default=16, help='Batch size')
    parser.add_argument('--num_epochs', type=int, default=8, help='Number of epochs to train')
    parser.add_argument('--log_dir', type=str, default='./logs', help='Directory to save logs')
    parser.add_argument('--save_dir', type=str, default='./saved_models', help='Directory to save trained models')
    
    # argument 파싱
    args = parser.parse_args()
    return args

 

args.parameter과 같은 방식으로 입력받은 argument를 사용할 수 있다.

# train.py

args = get_args()

data_dir = args.data_dir
model_type = args.model_type
lr = args.lr
batch_size = args.batch_size
num_epochs = args.num_epochs
log_dir = args.log_dir
save_dir = args.save_dir

 

CLI에서는 파이썬 스크립트 실행 시 아래와 같은 방식으로 argument의 값을 지정할 수 있다.

python train.py --data_dir ./dataset/path --model_type bert --lr 0.001 --batch_size 64 --num_epochs 20

 

omegaconf

python argparse를 써서 configuration을 관리하게 되면 실험을 진행할 때마다 수많은 옵션들을 CLI에 직접 입력해야하는 점이 번거롭기도 하고, 실수도 자주 발생한다. 이러한 문제를 해결하기 위해 omegaconf를 활용하여 YAML 파일에서 configuration 정보들을 관리할 수 있다.

아래는 configuration 정보를 담은 YAML 파일의 예시이다.

# config/default.yaml

model:
    model_type: bert

train:
    lr: 0.0001
    batch_size: 16
    num_epoch: 8

path:
    data_dir: ./dataset/path
    log_dir: ./logs
    save_dir: ./saved_models

 

YAML 파일에서 configuration 정보를 가져오는 방식은 아래와 같다.

1. argparse로 YAML 파일의 경로를 받아온다.

python train.py --config=config/default.yaml

2. omegaconf로 YAML 파일에 저장된 configuration 정보를 읽어온다.

import argparse

def get_args():
    parser = argparse.ArgumentParser(description='ML Experiment Configuration')
    
    parser.add_argument('--config', type=str, default='config/default.yaml')
    
    args = parser.parse_args()

    # YAML 파일을 읽어온다
    config = omegaconf.OmegaConf.load(args.config)
    return config

 

YAML 파일에서 작성된 계층 구조대로 configuration 값에 접근할 수 있다.

접근 방식은 dict 방식과 attribute 방식 모두 사용할 수 있다.

config = get_args()

print(config["model"]["model_type"] # dict 방식
>> bert 
print(config.train.lr) # attribute 방식
>> 0.0001

 

YAML 파일로 실험 configuration을 관리할 때의 장점은

• YAML 파일을 공유하여 손쉽게 실험을 재현할 수 있다.
• 서로 다른 두 실험의 configuration 차이를 한눈에 파악하기 쉽다.