Pytorch, 이미지 분류 코드 자세히 이해하기 (6편) - Training/valid/test 下편

2023.12.10 - [Deep Learning/Pytorch] - Pytorch, 이미지 분류 코드 자세히 이해하기 (5편) - Training/valid/test 上편

Pytorch, 이미지 분류 코드 자세히 이해하기 (5편) - Training/valid/test 上편

본격적으로 위 게시글에서 만든 함수를 사용하여 모델 학습 및 검증을 진행한다.

model = ResNet9(3, 30)
model = model.to('cuda')

all_df = df.reset_index(drop=True)
all_dataset = CustomDataset(all_df['img_path'], all_df['class'], train_aug)
all_loader = DataLoader(all_dataset, batch_size = CFG.batch_size, shuffle=True, num_workers=8)

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=CFG.lr)
scheduler = optim.lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lambda epoch: 0.95 ** epoch)

best_loss = np.Inf
train_depth_loss = torch.zeros(CFG.epochs)
valid_depth_loss = torch.zeros(CFG.epochs)
history_file = open("new_oversample_resnet9_noCutout.txt", "w")
bar = tqdm(range(0, CFG.epochs))
for e in bar:
    loss, acc = run_model(model, all_loader, criterion, optimizer, is_training=True)
    print('Epoch: {} \tTraining Loss: {:.6f}'.format(e, loss))
    print("lr: ", optimizer.param_groups[0]['lr'])
    scheduler.step()

    torch.save(model.state_dict(), f'{path}/final_model.pt')
history_file.close()

 

위 코드를 차근히 살펴보고 사용된 내용에 대한 이론지식을 정리한다.

* 아, 한 가지 강조할 점이 있다. 지금까지의 내용들은 교내 AI경진대회에서 진행한 내용이고, 기간이 일주일만 주어졌기 때문에 초반에는 K-fold 검증을 통해 대충 몇 에포크에서 수렴이 되는지 확인을 한 후, 시간 절약을 위해 전체 데이터를 검증 없이 한 번에 학습데이터로 사용했다. 따라서 위 코드에서 검증과정은 없을 것이다. run_model()의 인자로 is_training=False로 지정하고, validation 셋을 이용해 val_loader를 만들고 이를 인자로 all_loader 넘겨주면 검증과정이 된다.

model = ResNet9(3, 30)
model = model.to('cuda')

all_df = df.reset_index(drop=True)
all_dataset = CustomDataset(all_df['img_path'], all_df['class'], train_aug)
all_loader = DataLoader(all_dataset, batch_size = CFG.batch_size, shuffle=True, num_workers=8)

4편에 소개한 ResNet 모델을 사용하였고 이를 GPU장치로 옮겼다. all_df는 위 주의사항대로 모든 데이터를 학습으로 사용하고자 만들었다. 이때, reset_index를 통해 인덱싱 에러가 안나도로고 하였다. 2편에서 설명한 CustomDataset을 사용해 데이터셋을 만들고 이를 DataLoader에 전달한다. 배치사이즈는 실험적으로 128로 지정하였다.

 

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=CFG.lr)
scheduler = optim.lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lambda epoch: 0.95 ** epoch)

criterion은 CE(CrossEntropy)Loss를 사용한다. 이미지 다중분류 테스크에서 굉장히 많이 이용되는 loss이다. 이는 예측된 확률 분포와 실제 분포 사이의 차이를 효과적으로 측정한다. 실제 레이블과 예측값이 크게 다를 때 더욱 큰 손실을 부여하여 모델이 더 빠르고 효과적으로 학습하도록 돕는다.

optimizer은 AdamW를 사용했다. AdamW는 Adam의 weight decay의 방식을 개선한 것이 주된 특징이다. 먼저, Adam은 모멘텀(momentum)과 RMSprop의 장점을 결합하였다. Adam에서는 weight decay가 gradient 업데이트 단계에 포함되어있었지만, AdamW는 이를 분리하여 각 업데이트 단계 후에 독립적으로 적용한다.

scheduler는 LambdaLR을 사용했다. 주어진 람다 함수에 따라 learning rate를 조절한다. 여기서 말하는 람다함수는 파이썬 lambda 문법을 의미한다. 위 코드 기준으로 0.95^(epoch)로 설정해두었는데,

• epoch = 0일 때, 0.95 ** 0 = 1.0 (학습률은 초기값을 유지)
• epoch = 1일 때, 0.95 ** 1 = 0.95 (학습률이 초기값의 95%로 감소)
• epoch = 2일 때, 0.95 ** 2 ≈ 0.9025 (학습률이 초기값의 약 90.25%로 감소)

이런식으로 점차 learning rate를 낮추어간다. scheduler는 학습 초기에 높은 학습률로 빠르게 학습하고, 점차 낮추어가며 수렴을 돕는 전략으로 사용된다.