3. Brain Tumor 예측 인공지능 - 모델 훈련시키기

(라이브러리 버전은 이전 글을 확인해주세요.)

 

사용자가 선택한 값을 받아 훈련을 진행해야 합니다. 사용한 모델 중 가장 유명한 UNet 모델을 사용한 코드를 이용해 설명하겠습니다.

 

1. 사용 모델 : UNet

2. import 내용

from werkzeug.utils import secure_filename # 간단한 웹에 구현하기 위해 사용했습니다.

from flask import Flask, render_template, request, jsonify, render_template # Flask 을 이용해서 웹을 구현했습니다.

import torch # tensorflow 보단 pytorch 를 이용해 개발했습니다.

import warnings # 개발 IDE 가 jupyter notebook 이라 경고 메시지를 무시하기 위해 사용했습니다.

warnings.filterwarnings("ignore")

import os

import json

import numpy as np

import torch

import torch.nn as nn

import torch.optim as optim

from torch.utils.data import Dataset, DataLoader, random_split

import matplotlib.pyplot as plt

import cv2

from tqdm import tqdm

import matplotlib.pyplot as plt

import time

from sklearn.metrics import accuracy_score

from torchmetrics.classification import BinaryAccuracy

import torchvision

import torchvision.transforms as transforms

from torch.optim.lr_scheduler import ReduceLROnPlateau

import torchvision.models as models

import import_ipynb # jupyter notebook 에서 .ipynb 파일을 import 하기 위해 사용해야합니다.

import models_web # UNet 모델을 가지고 오기 위해 만든 파일입니다.

 

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # GPU 로 연산하기 위해 사용했습니다.

가장 중요한 점은 models_web.ipynb 파일에 UNet 모델을 구현한 코드가 있어야 합니다. 구조는 구글에 검색해서 얻은 구현 코드를 그대로 복붙 했습니다.

 

3. CustomDataset 과 DataLoader

class CustomDataset:

    def __init__(self, dataset):

        self.dataset = dataset

    def __len__(self):

        return len(self.dataset)

    def __getitem__(self, idx):

        image_path = self.dataset[idx]["image_path"]

        mask_path = self.dataset[idx]["mask_path"]

        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        _, mask = cv2.threshold(mask, 128, 1, cv2.THRESH_BINARY)

        image = torch.tensor(image, dtype=torch.float).unsqueeze(0)

        mask = torch.tensor(mask, dtype=torch.long)

        return image, mask

CustomDataset 을 구현해 주었습니다. 학습용 데이터를 그대로 사용할 수 없기 때문에 이미지를 변환시켜주는 작업이 필요합니다.

 

    dataset = CustomDataset(dataset)

    dataset_size = len(dataset)

    train_size = int(dataset_size * 0.8) # dataset 의 양이 적기 때문에 0.8 로 설정

    test_size = int(dataset_size * 0.1)

    val_size = dataset_size - train_size - test_size

 

    train_data, test_data, val_data = random_split(dataset, [train_size, test_size, val_size])

이렇게 해주면 학습하기 위해 사용되는 모델에 들어갈 수 있습니다.

    # 웹에서 입력받은 batch_size 값을 입력받아 설정

    Batch_Size = int(request.form['Batch_Size'])

    train_loader = DataLoader(train_data, batch_size=Batch_Size, shuffle=True)

    test_loader = DataLoader(test_data, batch_size=Batch_Size, shuffle=True)

    val_loader = DataLoader(val_data, batch_size=Batch_Size, shuffle=True)

항상 기억해야할 점은 학습용 데이터는 CustomDataset 으로 값을 변환 후 DataLoader 로 가져와야합니다.

 

4. 모델 학습  

    try:

        total_epoch_bar = tqdm(range(start_epoch, num_epochs), desc="Total progress", position=0, leave=False)

        for epoch in total_epoch_bar:

            tqdm_loader = tqdm(enumerate(train_loader), total=len(train_loader), desc=f"Epoch {epoch + 1}", position=1, leave=False)

            start_time = time.time()

            loss_values.append([])

            accuracy_values.append([])

            for i, (images, labels) in tqdm_loader:

                images = images.to(device)

                labels = labels.to(device)

                outputs = model(images)

                loss = criterion(outputs, labels)

                acc = accuracy(outputs, labels, images)

                optimizer.zero_grad()

                loss.backward()

                optimizer.step()

                loss_values[-1].append(loss.item())

                accuracy_values[-1].append(acc)

                tqdm_loader.set_postfix({"Loss": loss.item()})

                time_elapsed = time.time() - start_time

                estimated_time_remaining = (len(train_loader) - i - 1) * (time_elapsed / (i + 1))

                tqdm_loader.set_description(f"Epoch {epoch + 1} Accuracy {acc:.5f} (Estimated time remaining: {estimated_time_remaining:.2f}s)")

            loss_values_mean_each_epoch.append(mean(loss_values[-1]))

            accuracy_value_mean_each_epoch.append(mean(accuracy_values[-1]))

            training_visualization(outputs, labels, images)

            current_loss = loss_values_mean_each_epoch[-1]

            current_accuracy = accuracy_value_mean_each_epoch[-1]

            validation_loss_value.append([])

            validation_accuracy_value.append([])

           

            model.eval()

            total_loss = 0.0

            with torch.no_grad():

                for images, labels in val_loader:

                    images = images.to(device)

                    labels = labels.to(device)

                   

                    outputs = model(images)

                    validation_loss = criterion(outputs, labels)

                    total_loss += validation_loss.item()

                    validation_accuracy = accuracy(outputs, labels, images)

                   

                    validation_loss_value[-1].append(validation_loss.item())

                    validation_accuracy_value[-1].append(validation_accuracy)

                   

                    validation_loss_values_mean_each_epoch.append(mean(validation_loss_value[-1]))

                    validation_accuracy_value_mean_each_epoch.append(mean(validation_accuracy_value[-1]))

                   

            '''

            early stop 을 적용하기 위해 어느정도에서 훈련을 멈출지 값을 정해준다.

            학습이 조기 종료 되는 시점은 현재 Loss < 가장 낮은 Loss 값과 현재 Accuracy > 가장 높은 Accuracy 값인 경우

            아래의 조건을 참고해 조기종료를 실시한다.

            '''

            if current_loss < best_loss and current_accuracy > best_accuracy:

                average_loss = total_loss / len(val_loader)

                best_loss = current_loss

                best_accuracy = current_accuracy

                epochs_without_improvement = 0

                torch.save(model.state_dict(), model_save_path)

                print("average loss : ", average_loss)

                print(", validation loss mean each epoch : ", validation_loss_values_mean_each_epoch)

                print(", validation accuracy mean each epoch :  ", validation_accuracy_value_mean_each_epoch)

            else:

                epochs_without_improvement += 1

            if epochs_without_improvement >= patience:

                print(f'Early stopping! No improvement for {patience} epochs!')

                break

tqdm 으로 모델이 학습할 때 어느정도 진행되었는지 확인 할 수 있게 해줍니다.

 

<궁금증>

 

Q1.

    dataset = CustomDataset(dataset)

    dataset_size = len(dataset)

    train_size = int(dataset_size * 0.8) # dataset 의 양이 적기 때문에 0.8 로 설정

    test_size = int(dataset_size * 0.1)

    val_size = dataset_size - train_size - test_size

이 값은 얼마큼 해줘야 하나요?

 

A1. 말 그대로 CustomDataset 입니다. 자신이 원하는 값을 넣어주면 됩니다. 훈련과 테스트 비율을 7:3 으로 해도 되고, 저는 훈련에 사용한 데이터셋의 양이 적어 8:1:1 로 진행했습니다.

 

Q2. train, test, val 의 의미는 뭔가요?

A2.

 

Train data(학습 데이터):

- Train data 는 모델을 학습하는데 사용되는 데이터셋 입니다.

- 모델에 넣어주기 위해 데이터로부터 특징(feature) 과 레이블(label) 을 추출하고, 이 값을 넣어주면 패턴을 학습합니다.

 

Test data(테스트 데이터):

- Test data 는 학습된 모델의 성능을 평가하는데 사용되는 데이터셋 입니다.

- 마치 기출문제집으로 공부한 내용을 테스트 하는것과 같습니다.

 

Validation data(검증 데이터):

- 모델이 데이터를 학습하는 중간중간 성능을 평가하기 위해 사용됩니다.

-  모델은 데이터만으로 학습하면 좋은 값을 얻을 수 없습니다. 따라서 모델을 학습할 때 여러 옵션을 설정해주는데, 이런 옵션을 이용해 학습을 최적화 합니다.

 

Q3. 훈련은 어떤 방법으로 진행하나요?

A3. 학습 데이터를 이용해 모델이 데이터를 학습 - 검증 데이터를 이용해 학습하며 중간중간 알맞게 학습하나 확인 - 학습이 끝나고 얼마나 정확하게 학습했나 테스트 데이터를 이용해 비교

순으로 진행됩니다.

 

(전체 코드는 차후 소프트웨어 등록이 끝나면 주석과 함께 첨부할 예정입니다.)