Link

실습: train.py 구현하기

사용자는 train.py를 통해 다양한 파라미터를 시도하고 모델을 학습할 수 있습니다. CLI 환경에서 바로 train.py를 호출 할 것이며, 그럼 train.py의 다음 코드가 실행될 것입니다.

if __name__ == '__main__':
    config = define_argparser()
    main(config)

먼저 define_argparser라는 함수를 통해 사용자가 입력한 파라미터들을 config라는 객체에 저장합니다. 다음 코드는 define_argparser 함수를 정의한 코드입니다.

def define_argparser():
    p = argparse.ArgumentParser()

    p.add_argument('--model_fn', required=True)
    p.add_argument('--gpu_id', type=int, default=0 if torch.cuda.is_available() else -1)

    p.add_argument('--train_ratio', type=float, default=.8)

    p.add_argument('--batch_size', type=int, default=256)
    p.add_argument('--n_epochs', type=int, default=20)

    p.add_argument('--n_layers', type=int, default=5)
    p.add_argument('--use_dropout', action='store_true')
    p.add_argument('--dropout_p', type=float, default=.3)

    p.add_argument('--verbose', type=int, default=1)

    config = p.parse_args()

    return config

argparse 라이브러리를 통해 다양한 입력 파라미터들을 손쉽게 정의하고 처리할 수 있습니다. train.py와 함께 주어질 수 있는 입력들은 다음과 같습니다.

파라미터 이름 설명 디폴트 값
model_fn 모델 가중치가 저장될 파일 경로 없음. 사용자 입력 필수
gpu_id 학습이 수행될 그래픽카드 인덱스 번호 (0부터 시작) 0 또는 그래픽 부재시 -1
train_ratio 학습 데이터 내에서 검증 데이터가 차지할 비율 0.8
batch_size 미니배치 크기 256
n_epochs 에포크 갯수 20
n_layers 모델의 계층 갯수 5
use_dropout 드랍아웃 사용여부 False
dropout_p 드랍아웃 사용시 드랍 확률 0.3
verbose 학습시 로그 출력의 정도 1

model_fn 파라미터는 required=True 가 되어 있으므로 실행시 필수적으로 입력되어야 합니다. 이외에는 디폴트 값이 정해져 있으므로, 사용자가 따로 지정해주지 않으면 디폴트 값이 적용됩니다. 만약 다른 알고리즘의 도입으로 이외에도 추가적인 하이퍼파라미터의 설정이 필요하다면 add_argument 함수를 통해 프로그램이 입력을 받도록 설정할 수 있습니다. 이렇게 입력받은 파라미터들은 다음과 같이 접근할 수 있습니다.

config.model_fn

앞서 모델 클래스를 정의할 때, hidden_sizes라는 리스트를 통해 쌓을 블럭들의 크기들을 지정할 수 있었습니다. 사용자가 블럭 크기들을 일일히 지정하는것은 어쩌면 번거로운 일이 될 수 있기 때문에, 사용자가 모델의 계층 갯수만 정해주면 자동으로 등차수열을 적용하여 hidden_sizes를 구하고자 합니다. 다음의 get_hidden_sizes 함수는 해당 작업을 수행합니다.

def get_hidden_sizes(input_size, output_size, n_layers):
    step_size = int((input_size - output_size) / n_layers)

    hidden_sizes = []
    current_size = input_size
    for i in range(n_layers - 1):
        hidden_sizes += [current_size - step_size]
        current_size = hidden_sizes[-1]

    return hidden_sizes

이제 학습에 필요한 대부분의 부분들이 구현되었습니다. 이것들을 모아서 학습을 진행하도록 코드를 구현하면 됩니다. 다음의 코드는 앞서 구현한 코드들을 모아서 실제 학습을 진행하는 과정을 수행하도록 구현한 코드입니다.

def main(config):
    # Set device based on user defined configuration.
    device = torch.device('cpu') if config.gpu_id < 0 else torch.device('cuda:%d' % config.gpu_id)

    # Load data and split into train/valid dataset.
    x, y = load_mnist(is_train=True, flatten=True)
    x, y = split_data(x.to(device), y.to(device), train_ratio=config.train_ratio)

    print("Train:", x[0].shape, y[0].shape)
    print("Valid:", x[1].shape, y[1].shape)

    # Get input/output size to build model for any dataset.
    input_size = int(x[0].shape[-1])
    output_size = int(max(y[0])) + 1

    # Build model using given configuration.
    model = ImageClassifier(
        input_size=input_size,
        output_size=output_size,
        hidden_sizes=get_hidden_sizes(input_size,
                                      output_size,
                                      config.n_layers),
        use_batch_norm=not config.use_dropout,
        dropout_p=config.dropout_p,
    ).to(device)
    optimizer = optim.Adam(model.parameters())
    crit = nn.NLLLoss()

    if config.verbose >= 1:
        print(model)
        print(optimizer)
        print(crit)

    # Initialize trainer object.
    trainer = Trainer(model, optimizer, crit)

    # Start train with given dataset and configuration.
    trainer.train(
        train_data=(x[0], y[0]),
        valid_data=(x[1], y[1]),
        config=config
    )

    # Save best model weights.
    torch.save({
        'model': trainer.model.state_dict(),
        'opt': optimizer.state_dict(),
        'config': config,
    }, config.model_fn)

MNIST에 특화된 입출력 크기를 갖는 것이 아닌, 벡터 형태의 어떤 데이터도 입력을 받아 분류할 수 있도록 input_size와 output_size 변수를 계산하는 것에 주목하세요. MNIST에 특화된 하드코딩을 제거하였기 때문에, load_mnist 함수가 아닌 다른 로딩 함수로 바꿔치기 한다면 이 코드는 얼마든지 바로 동작할 수 있습니다. 사용자로부터 입력받은 설정configuration을 활용하여 모델을 선언한 이후에, Adam 옵티마이저와 NLL 손실 함수도 함께 준비합니다. 그리고 트레이너를 초기화한 후, train함수를 호출하여 불러온 데이터를 넣어주어 학습을 시작합니다. 학습이 종료된 이후에는 torch.save 함수를 활용하여 모델 가중치를 config.model_fn 경로에 저장합니다.

코드 실행하기

이번 프로젝트는 주피터 노트북이 아닌 파이썬 스크립트로 작성되었기 때문에, CLI 환경(e.g. DOS 커맨드 창)에서 실행할 수 있습니다. train.py에서 argparse 라이브러리를 활용하여 사용자의 입력을 파싱하여 인식할 수 있죠. 다만, 처음 사용하거나 오랜만에 실행하는 경우, 어떤 입력 파라미터들이 가능한지 잘 기억이 안날 수도 있습니다. 그땐 입력 파라미터 없이 다음과 같이 실행하거나 ‘–help’ 파라미터를 넣어 실행하면 입력 가능한 파라미터들을 확인할 수 있습니다.

$ python train.py
usage: train.py [-h] --model_fn MODEL_FN [--gpu_id GPU_ID] [--train_ratio TRAIN_RATIO] [--batch_size BATCH_SIZE] [--n_epochs N_EPOCHS] [--n_layers N_LAYERS] [--use_dropout] [--dropout_p DROPOUT_P] [--verbose VERBOSE]
train.py: error: the following arguments are required: --model_fn

입력할 파라미터들을 정해서 다음과 같이 직접 실행하게 되면 학습이 정상적으로 진행되는 것을 볼 수 있습니다.

$ python train.py --model_fn tmp.pth --gpu_id -1 --batch_size 256 --n_epochs 20 --n_layers 5
Train: torch.Size([48000, 784]) torch.Size([48000])
Valid: torch.Size([12000, 784]) torch.Size([12000])
ImageClassifier(
  (layers): Sequential(
    (0): Block(
      (block): Sequential(
        (0): Linear(in_features=784, out_features=630, bias=True)
        (1): LeakyReLU(negative_slope=0.01)
        (2): BatchNorm1d(630, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): Block(
      (block): Sequential(
        (0): Linear(in_features=630, out_features=476, bias=True)
        (1): LeakyReLU(negative_slope=0.01)
        (2): BatchNorm1d(476, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (2): Block(
      (block): Sequential(
        (0): Linear(in_features=476, out_features=322, bias=True)
        (1): LeakyReLU(negative_slope=0.01)
        (2): BatchNorm1d(322, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (3): Block(
      (block): Sequential(
        (0): Linear(in_features=322, out_features=168, bias=True)
        (1): LeakyReLU(negative_slope=0.01)
        (2): BatchNorm1d(168, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (4): Linear(in_features=168, out_features=10, bias=True)
    (5): LogSoftmax(dim=-1)
  )
)
Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    eps: 1e-08
    lr: 0.001
    weight_decay: 0
)
NLLLoss()
Epoch(1/20): train_loss=1.9761e-01  valid_loss=9.7320e-02  lowest_loss=9.7320e-02
Epoch(2/20): train_loss=8.0091e-02  valid_loss=9.0333e-02  lowest_loss=9.0333e-02
Epoch(3/20): train_loss=5.6217e-02  valid_loss=8.5263e-02  lowest_loss=8.5263e-02
Epoch(4/20): train_loss=4.2374e-02  valid_loss=7.9394e-02  lowest_loss=7.9394e-02
Epoch(5/20): train_loss=3.0406e-02  valid_loss=8.6583e-02  lowest_loss=7.9394e-02
Epoch(6/20): train_loss=2.9545e-02  valid_loss=7.4341e-02  lowest_loss=7.4341e-02
Epoch(7/20): train_loss=2.1318e-02  valid_loss=7.6004e-02  lowest_loss=7.4341e-02
Epoch(8/20): train_loss=2.2078e-02  valid_loss=7.5559e-02  lowest_loss=7.4341e-02
Epoch(9/20): train_loss=1.6825e-02  valid_loss=7.3756e-02  lowest_loss=7.3756e-02
Epoch(10/20): train_loss=1.6661e-02  valid_loss=8.4252e-02  lowest_loss=7.3756e-02
Epoch(11/20): train_loss=1.9243e-02  valid_loss=7.6462e-02  lowest_loss=7.3756e-02
Epoch(12/20): train_loss=1.3275e-02  valid_loss=7.4506e-02  lowest_loss=7.3756e-02
Epoch(13/20): train_loss=1.2972e-02  valid_loss=7.4413e-02  lowest_loss=7.3756e-02
Epoch(14/20): train_loss=1.3329e-02  valid_loss=8.3443e-02  lowest_loss=7.3756e-02
Epoch(15/20): train_loss=1.0855e-02  valid_loss=8.3068e-02  lowest_loss=7.3756e-02
Epoch(16/20): train_loss=1.0555e-02  valid_loss=6.7707e-02  lowest_loss=6.7707e-02
Epoch(17/20): train_loss=7.2489e-03  valid_loss=7.7850e-02  lowest_loss=6.7707e-02
Epoch(18/20): train_loss=7.0233e-03  valid_loss=8.2895e-02  lowest_loss=6.7707e-02
Epoch(19/20): train_loss=1.2123e-02  valid_loss=9.4855e-02  lowest_loss=6.7707e-02
Epoch(20/20): train_loss=1.2536e-02  valid_loss=1.0303e-01  lowest_loss=6.7707e-02

만약 주피터 노트북을 활용하여 실행해야 했다면, 보통은 코드에 하이퍼파라미터들을 직접 하드코딩 한 후, 실험 때마다 필요한 값으로 일일히 바꿔서 실행해야 했을 것입니다. 그럼 사용자가 필요한 값을 바꿔야 할 때 번거로운 것을 떠나서도, 실수로 깜빡하고 바꾸지 않는 등 여러가지 문제가 발생할 수 있습니다. 하지만 앞서 본 것처럼 CLI 환경에서 사용자가 train.py를 실행할 때 간단하게 입력으로 필요한 파라미터를 주도록하면 간편하고 좀 더 실수의 여지도 줄어들 수 있을 것입니다. MNIST 분류기의 성능을 높이는 것도 중요하지만, 이처럼 MNIST 분류기를 좀 더 잘 연구하기 위해 편리한 실험 환경을 갖추는 것도 매우 중요합니다.