[My IT : Codes] U-Net 활용 Sementic Segmentation : Football Dataset(2) (하이퍼파라미터 튜닝 ~ 결론)

이전 코드 : https://uj07096.tistory.com/65

[My IT : Codes] U-Net 활용 Sementic Segmentation : Football Dataset(1) (시작~ 모델링)

 

5. 하이퍼파라미터 튜닝

모든 모델에 대해 하이퍼파라미터 튜닝을 진행하면 좋겠지만, 시간 효율성상 가장 빠르게 구동이 가능한 Custom U-Net으로 하이퍼파라미터 튜닝을 진행 후, 같은 하이퍼 파라미터로 다른 모델들의 학습도 진행하기로 하였다.

 

#optuna 활용 하이퍼파라미터 튜닝
def objective(trial) :
  learning_rate = trial.suggest_float('learning_rate', 1e-5, 1e-1, log = True)
  weight_decay = trial.suggest_float('weight_decay', 1e-5, 1e-2, log = True)

  #모델 할당
  model = UNet().to(device)

  #optimizer(AdamW) 설정
  optimizer = torch.optim.AdamW(model.parameters(), lr = learning_rate, weight_decay = weight_decay)

  #손실함수 설정
  criterion_focal = FocalLoss(gamma=2.0) #focal loss
  criterion_dice = DiceLoss() # 앞서 정의한 DiceLoss 클래스

  #튜닝 기준 지표
  dice = DiceScore(num_classes = 11, input_format = 'mixed').to(device)

  best_dice = 0.0
  epochs = 20

  #튜닝 시작
  for epoch in range(epochs) :
    model.train()
    for batch_idx, (images, masks) in enumerate(train_loader):
      images = images.to(device)
      masks = masks.to(device)

      #forward
      outputs = model(images)

      loss_focal = criterion_focal(outputs, masks)
      loss_dice = criterion_dice(outputs, masks)
      # loss 계산(스케일이 크게 다른 경우를 대비해 0.5해서 합침)
      loss = 0.5 * loss_focal + 0.5 * loss_dice

      #backward
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

    model.eval()
    dice.reset()
    with torch.no_grad():
      for images, masks in val_loader :
        images = images.to(device)
        masks = masks.to(device)

        outputs = model(images)

        dice.update(outputs, masks)

    dice_score = dice.compute().item()

    #best dice score 갱신시
    if dice_score > best_dice :
      best_dice = dice_score

    #pruning 추
    trial.report(dice_score, epoch)

    if trial.should_prune() :
      raise optuna.exceptions.TrialPruned()
  return best_dice


study = optuna.create_study(direction = 'maximize',
                            pruner = optuna.pruners.MedianPruner(n_startup_trials = 5,
                                                                 n_warmup_steps = 5))
study.optimize(objective, n_trials = 30)

print(f'Best Dice-Score:, {study.best_value:.5f}')
print('Best Params:', study.best_params)

 

 

---

6. 모델 학습/ 결과 시각화

6-1. 반복 함수 정의

총 3개의 모델을 학습하기 위해, 반복적으로 호출할 작업들을 함수로 정의하였다.

- train_test : 모델 학습 및 평가지표 dict에 저장 함수
- plot_history : train_loss와 val_loss의 그래프 시각화, dice score과 miou의 시각화
- draw_res : best model을 가져와 validation 데이터셋의 원본 이미지, 라벨 이미지, 예측 이미지를 시각화

 

#모델 학습 함수
def train_test(model, model_name, epochs, lr, wd) :
  model = model.to(device)

  #optimizer(AdamW) 설정
  optimizer = torch.optim.AdamW(model.parameters(), lr = lr, weight_decay = wd)

  #스케줄러 설정
  scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='max', patience=5, factor=0.5
)

  #손실함수 설정
  criterion_focal = FocalLoss(gamma=2.0) #focal loss
  criterion_dice = DiceLoss() # 앞서 정의한 DiceLoss 클래스

  #튜닝 기준 지표
  dice = DiceScore(num_classes = 11, input_format = 'mixed').to(device)
  miou = MeanIoU(num_classes = 11, input_format='mixed').to(device)

  #학습 기록 저장할 dict
  history = {'train_loss': [], 'val_loss' : [], 'val_miou': [], 'val_dice': [], 'time' : 0.0}

  #체크포인트 저장 폴더 생성
  os.makedirs('checkpoints', exist_ok=True)

  best_dice = 0.0
  best_epoch = 0
  stop_patience = 15
  stopping_cnt = 0

  print('학습 시작')
  start_time = time.time()
  for epoch in range(epochs) :
    model.train()
    t_loss = 0
    for batch_idx, (images, masks) in enumerate(train_loader):
      images = images.to(device)
      masks = masks.to(device)

      #forward
      outputs = model(images)

      loss_focal = criterion_focal(outputs, masks)
      loss_dice = criterion_dice(outputs, masks)
      # loss 계산(스케일이 크게 다른 경우를 대비해 0.5해서 합침)
      loss = 0.5 * loss_focal + 0.5 * loss_dice

      #backward
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      t_loss += loss.item()
      if (batch_idx + 1) % 4 == 0: print('.', end='')

    avg_t_loss = t_loss / len(train_loader)
    history['train_loss'].append(avg_t_loss)

    #Validation
    model.eval()
    dice.reset()
    miou.reset()
    v_loss = 0
    with torch.no_grad():
      for images, masks in val_loader :
        images = images.to(device)
        masks = masks.to(device)

        outputs = model(images)

        # loss 계산
        loss_focal = criterion_focal(outputs, masks)
        loss_dice = criterion_dice(outputs, masks)
        loss = 0.5 * loss_focal + 0.5 * loss_dice

        #평가지표 계산
        v_loss += loss.item()
        dice.update(outputs, masks)
        miou.update(outputs, masks)

    avg_v_loss = v_loss / len(val_loader)
    dice_score = dice.compute().item()
    miou_score = miou.compute().item()
    history['val_loss'].append(avg_v_loss)
    history['val_dice'].append(dice_score)
    history['val_miou'].append(miou_score)

    if (epoch + 1) % 5 == 0 :
      print(f"Epoch [{epoch+1}/{epochs}]")
      print(f"Train Loss: {avg_t_loss:.4f}")
      print(f"Val Dice: {dice_score:.4f}, MIoU : {miou_score:.4f}")

    scheduler.step(dice_score)
    if dice_score > best_dice :
      best_dice = dice_score
      best_epoch = epoch + 1
      torch.save(model.state_dict(), f'checkpoints/{model_name}_best_model.pt')
      print(f"Best 모델 갱신 (Epoch {best_epoch}, Dice Score : {best_dice:.6f})")
      stopping_cnt = 0 #stopping count 리셋

    else :
      stopping_cnt += 1
      if stopping_cnt >= stop_patience :  #Early Stopping
        print(f"Early Stopping (Epoch {epoch + 1})")
        break

  end_time = time.time()
  history['time'] = end_time - start_time

  print('학습 종료')
  print(f'모델 : {model_name}, 소요시간 : {end_time - start_time: .5f} (sec)')
  print(f"최적 Dice Score 에포크: {best_epoch}, 최고 Dice 스코어: {best_dice:.6f}")
  return history

 

loss 그래프의 경우, 1 ~ 20 epoch는 loss값이 커서 뒤쪽에서 과적합 여부를 확인하기 어려우므로 20 epoch 이후부터 그래프를 그렸다. miou와 dice score의 경우 1 epoch부터 관찰했다.

 

#history 시각화 함수 정의
def plot_history(history) :
  start = 20
  epochs_range_loss = range(start, len(history['train_loss']) + 1)
  epochs_range_score = range(1, len(history['val_miou']) + 1)  # ax2용

  plt.figure(figsize=(8, 8))

  # 학습&평가 손실 시각화
  ax1 = plt.subplot(2, 1, 1)
  # ax1 - 20 epoch 이후
  ax1.plot(epochs_range_loss, history['train_loss'][start-1:], label='train_loss')
  ax1.plot(epochs_range_loss, history['val_loss'][start-1:], label='val_loss')

  ax1.set_title('Training vs Val Loss')
  ax1.set_xlabel('Epochs')
  ax1.set_ylabel('Loss')
  ax1.legend()

  # 정확도&F1 지표 시각화
  ax2 = plt.subplot(2, 1, 2)
  # ax2 - 1부터 전체
  ax2.plot(epochs_range_score, history['val_miou'], label='val_miou')
  ax2.plot(epochs_range_score, history['val_dice'], label='val_dice')
  ax2.set_title('Validation Dice Score & Miou')
  ax2.set_xlabel('Epochs')
  ax2.set_ylabel('Score')
  ax2.legend()


  plt.legend()
  plt.tight_layout()
  plt.show()

 

# 예측 시각화 함수 정의
def draw_res(model, model_name, n = 5):
    #해당 best 모델 로드
    model = model
    model.load_state_dict(torch.load(f'checkpoints/{model_name}_best_model.pt'))
    model.to(device)
    model.eval()

    #왼쪽 : 오리지널 이미지, 가운데, : 예측 이미지, 오른쪽 : 라벨 이미지
    fig, axes = plt.subplots(n, 3, figsize=(15, n * 4))
    axes[0][0].set_title('Original Image', fontsize=12)
    axes[0][1].set_title('Prediction', fontsize=12)
    axes[0][2].set_title('Label Mask', fontsize=12)

    cnt = 0
    with torch.no_grad():
        for images, masks in val_loader:
            images = images.to(device)
            outputs = model(images)
            preds = torch.argmax(outputs, dim=1).cpu()  # (B, H, W)

            for i in range(images.size(0)):
                if cnt >= n:
                    break

                # 원본 이미지
                img = images[i].cpu().permute(1, 2, 0).numpy()
                img = np.clip(img, 0, 1)
                axes[cnt][0].imshow(img)
                axes[cnt][0].axis('off')

                # 예측 마스크
                pred_color = index_to_color(preds[i].numpy())
                axes[cnt][1].imshow(pred_color)
                axes[cnt][1].axis('off')

                # 정답 마스크
                label_color = index_to_color(masks[i].numpy())
                axes[cnt][2].imshow(label_color)
                axes[cnt][2].axis('off')

                cnt += 1

            if cnt >= n:
                break

    plt.tight_layout()
    plt.show()

---

6-2. Custom U-Net

모델 시각화

model = UNet().to(device)
summary(model, input_size=(1, 3, 512, 512), depth=1)  # depth로 깊이 조절

모델 학습 / 그래프 시각화

#튜닝에서 얻은 wd와 lr 로드, 모델 설정
wd = study.best_params['weight_decay']
lr = study.best_params['learning_rate']
epochs = 150

model_name = 'Unet'

#모델 학습
Unet_history = train_test(model, model_name, epochs = epochs, lr = lr, wd = wd)

 

 

plot_history(Unet_history)

 

예측 결과 시각화

---

6-3. ResNet BackBone(Untrained)

모델 시각화

model = UNetResNet34().to(device)

summary(model, input_size=(1, 3, 512, 512), depth=1)  # depth로 깊이 조절

 

모델 학습/ 그래프 시각화

model_name = 'Unet_ResNet_Untrained'

#모델 학습
Unet_ResNet_Untrained_history = train_test(model, model_name, epochs = epochs, lr = lr, wd = wd)

 

 

plot_history(Unet_ResNet_Untrained_history)

 

예측 결과 시각화

draw_res(model, model_name)

---

6-4. ResNet BackBone(PreTrained)

모델의 아키텍처는 UNetResNet34()와 같고, 사전학습 여부만 다르므로 시각화 결과는 6-3과 같다.

모델 학습 / 그래프 시각화

model = UNetResNet34(weights = ResNet34_Weights.IMAGENET1K_V1)
model_name = 'Unet_ResNet_PreTrained'

#모델 학습
Unet_ResNet_Pretrained_history = train_test(model, model_name, epochs = epochs, lr = lr, wd = wd)

plot_history(Unet_ResNet_Pretrained_history)

 

예측 결과 시각화

draw_res(model, model_name)

 

7. 성능지표 비교 / 결론

Dice Score이 가장 높았던 시점의 모델들(best_model.pt)의 당시의 성능지표를 DataFrame으로 생성해 표를 만들어 관찰하고, 시각화하였다.

 

DataFrame 시각화

all_history = {
    'Unet' : Unet_history,
    'Unet_ResNet_Untrained' : Unet_ResNet_Untrained_history,
    'Unet_ResNet_PreTrained' : Unet_ResNet_Pretrained_history
}

rows = []
for model_name, h in all_history.items():
    best_epoch = h['val_dice'].index(max(h['val_dice']))  # best dice score 찍은 에폭
    rows.append({
        'model': model_name,
        'best_dice': round(max(h['val_dice']), 4),
        'best_miou': round(h['val_miou'][best_epoch], 4),
        'train_loss': round(h['train_loss'][best_epoch], 4),
        'time (s)': round(h['time'], 1)
    })

df_result = pd.DataFrame(rows).set_index('model')
df_result

 

막대그래프 시각화

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
fig.suptitle('Model Comparison', fontsize = 20)

metrics = ['train_loss', 'time (s)', 'best_dice', 'best_miou']
titles = ['Train_Loss', 'Time (s)', 'Best Dice Score', 'Best Miou']
colors = ['steelblue', 'seagreen', 'tomato', 'mediumpurple']
formats = ['.4f', '.1f', '.4f', '.4f']


for ax, metric, title, color, fmt in zip(axes.flatten(), metrics, titles, colors, formats):
    bars = ax.bar(df_result.index, df_result[metric], color=color, alpha=0.8)
    ax.set_title(title, fontsize=13)
    ax.set_xticks(range(len(df_result.index)))
    ax.set_xticklabels(df_result.index, rotation=45, ha='right')
    ax.set_yscale('log')  #값 차이를 명확히 보기 위해서 logscale

    for bar, val in zip(bars, df_result[metric]):
        ax.text(bar.get_x() + bar.get_width()/2,
                bar.get_height() + 0.001,
                f'{val:{fmt}}',
                ha='center', va='bottom', fontsize=9)

plt.tight_layout()
plt.show()

 

---

8. 결론

backbone 유무 효과 비교

직접 구현한 U-Net 아키텍처가, BackBone을 Resnet으로 교체한 모델보다 시간이 훨씬 오래걸렸다. 왜 그랬을까 ?
-> ResNet 자체가 Pytorch 내부적으로 최적화가 잘 되어있고, 그때문에 같은 연산량이라도 더 빠르게 돌아간다.


UNet enc1: (B, 64, 512, 512)  ← 512 해상도에서 conv(풀링 X)
ResNet enc1: (B, 64, 256, 256) ← 바로 256으로 줄여버림

또한 이렇게 ResNet은 초반에 해상도를 확 줄여서 고해상도에서의 연산 자체가 적다.

 

사전학습된 weight 효과 비교

사전학습된 가중치를 가져왔을 때, 학습 시간 자체도 훨씬 빠르게 끝나고 정확도도 높게 나왔으며, 육안으로 예측 결과를 봤을 때에도 우수한 성능을 보였다. 왜 그랬을까 ?

 

아키텍처만 가지고 오고 사전학습된 weights가 없는 모델 : 114 epoch에서 수렴
사전학습된 weights를 가져온 모델 : 97 epoch에서 수렴

 

사전학습된 weights를 가져오게 되면, 처음부터 의미있는 features를 뽑기 때문에 빠르게 수렴하게 되고, 사전학습 되지 않은 모델을 가져오게 되면 가중치가 random이기 때문에 처음부터 feature 추출까지 진행을 해야한다. 따라서 더 적은 epochs으로도(적은 시간) 좋은 성능에 도달할 수 있었다.

 

 

 

Full code : Github

-> Sementic_Segmentation_Football_Dataset.ipynb

Deep_Learning_prac/Sementic_Segmentation_Football_Dataset.ipynb at main · jaeheonki/Deep_Learning_prac