1. 모델 분석/Segmentation

[2D Multi-Class Semantic Segmentation] 3. Model Training, Evaluation and Result

M_AI 2021. 5. 31. 21:05

 

안녕하세요! M_AI 입니다!

 

이전까지 데이터에 대한 설명과 학습을 위한 데이터 전처리 및 데이터셋 준비법을 설명했습니다!

 

이전글 : [2D Multi-Class Semantic Segmentation] 2. U-Net with Backbone and Loss Function

https://yhu0409.tistory.com/8

 

[2D Multi-Class Semantic Segmentation] 2. U-Net with Backbone and Loss Function

안녕하세요! M_AI 입니다! 이전글에서는 당뇨망막병증(Diabetic Retinopathy)에서 무엇을 세그멘테이션할 지와 이를 위해 데이터셋을 어떻게 준비하는 지에 대해 설명했습니다. 이번 장에서는 ImageNet

yhu0409.tistory.com

 


PC 화면으로 보는 것을 추천!!!!!

 

이 글을 읽으면 좋을 것 같은 사람들!

- 기본적인 분류(classfication) 모델을 Tensorflow 2.x 버전으로 구현할 줄 아는 사람

- Clssification을 넘어서 Object detection을 하길 원하는데 오픈 소스를 봐도 이해가 가질 않는 사람.

- 오픈 소스를 이해하더라도 입맛대로 수정하는 방법을 모르는 사람.

 

 

본 게시글의 목적

※ 본 게시글은 딥러닝 모델의 상세한 이론적 설명보다는 모델 구조를 바탕으로 코드 분석하고 이해하는 데에 목적에 있습니다!

※ 상세한 이론을 원하면 다른 곳에서 찾아보시길 바랍니다!

 

본 게시글의 특징

① Framework : Tensorflow 2.x

② 학습 환경

- Google Colab

③ Batch size = 1

 IDRiD Diabetic Retinopathy, DIARETDB1 데이터셋 사용하여 Fundus Image에서 병변 분할(Lesion Segmentation) 목적

( IDRiD 데이터셋 출처 : https://idrid.grand-challenge.org/Data/)

 

IDRiD - Grand Challenge

This challenge evaluates automated techniques for analysis of fundus photographs. We target segmentation of retinal lesions like exudates, microaneurysms, and hemorrhages and detection of the optic disc and fovea. Also, we seek grading of fundus images acc

idrid.grand-challenge.org

 

( DIARETDB1 데이터셋 출처 : https://www.it.lut.fi/project/imageret/diaretdb1/#DOWNLOAD)

 

DIARETDB1 - STANDARD DIABETIC RETINOPATHY DATABASE

DIARETDB1 - Standard Diabetic Retinopathy Database Calibration level 1 This is a public database for benchmarking diabetic retinopathy detection from digital images. The main objective of the design has been to unambiguously define a database and a testing

www.it.lut.fi

⑤ 게시글에서는 반말

 


 

3. Model Training, Evaluation and Result

3.1. Model Training and Evaluation

 

A. K-Folding Cross Validation

본 게시글에서는 모델 학습을 K-Fold Cross Validation으로 진행한다고 하였다. 그러한 이유는, 매우 부족한 데이터로는 모델 성능에 대한 신뢰성이 떨어지기에 이를 해결하고 K-Fold Cross Validation을 도입하였다.

 

처음에는 K-Fold를 하기 위해 scikit-learn의 Kfold 함수를 사용하려 했으나, 이 함수를 사용하려면 데이터셋 전체를 Load하여 split을 해야한다. 하지만 이 과정에 Out-of-Memory가 발생하였다.

 

이러한 이유는, 데이터 수는 적지만 데이터 증식(data augmentation)을 하면 용량이 매우 커져 OOM 발생한다.

 

이를 위해서 RAM에 올리지 않고, generator를 사용하면 data augmentation도 하고, K-Fold를 하는 방법을 찾아야만 했다. 이에 대한 해결법은 다음과 같다.

 

본 게시글에서 Training / Test dataset 수는 100개 / 36개이다. 즉, Training dataset 100개를 K개의 fold로 split해야하는데, 여기서는 5-Fold로 진행하였다.

 

 

(지난 글 [2D Multi-Class Semantic Segmentation] 1. Data Preprocessing의 1.2.1. Member Method에서 3. train_generator():과 4. valid_generator(): 코드를 보면서 이 글을 읽기 바란다.)

https://yhu0409.tistory.com/6

 

[2D Multi-Class Semantic Segmentation] 1. Data Preprocessing

2021.05.11 (화) 수정 내역 1. 데이터 추가 - 검증 데이터가 적어 모델에 대한 신뢰성이 떨어져, DIARETDB1 데이터셋을 추가. 2. 클래스 변경 - 기존에는 optic disc를 추가하여 함께 segmentation을 했으나 새.

yhu0409.tistory.com

 

K Validation Dataset (1개 Fold) Training Dataset (4개 Fold)
1 0 <= index < 20 20 <= index < 100
2 20 <= index < 40 0 <= index < 20 or 40 <= index < 100
3 40 <= index < 60 0 <= index < 40 or 60 <= index < 100
4 60 <= index < 80 0 <= index < 60 or 80 <= index < 100
5 80 <= index < 100 0 <= index < 80

 

위의 table은, 100개의 Training data에서 K값에 따라 특정 index에 있는 데이터를 validation과 training dataset으로 결정하도록 하는 방식이다.

 

그로 인해, train_generator에서 index가 범위 밖의 숫자면 그 데이터는 validation이므로 무시하고, 범위 내에 있으면 해당 데이터를 전처리하고 yield 한다.

valid_generator에서도 마찬가지고 index가 범위 밖이면 training이므로 무시하고, 범위 내에 있으면 해당 데이터를 전처리하고 yield한다.

 

아래는 K에 따라 validation과 training data를 결정하고 generator하는 Pseudo code이므로 참고하길 바란다.

# Pseudo Code
# Dataset의 길이는 100

Shuffle(Dataset) # 학습 전 데이터셋 한 번 셔플
for k = 1 to 5: # K-Fold 반복문
    for i = 1 to 100: # 데이터셋 반복문

        # Validataion data Range
        if 20*(k-1) <= i < 20*k:
            Dataset[i] is validation data
            yield Dataset[i]
            yield DataAugmentation(Dataset[i])

        # Training data Range
        else:
            Dataset[i] is training data
            yield Dataset[i]
            yield DataAugmentation(Dataset[i])

 

이런식으로 진행하면 RAM의 Out-Of-Memory 문제를 해결할 수 있다.

 

하지만 문제가 하나 있다면, 학습 도중마다 generator에서 데이터 증식(Data augmentation) 과정을 해야하므로 학습 시간이 조금 더 오래 걸릴 수 있다. 데이터 증식 과정은 CPU로 진행되는데, CPU 성능도 뛰어나다면 좋겠지만, Google Colab Pro에서는 CPU 성능이 그리 좋지 않기에 매우 아쉬운 부분이다. (그래도 GPU가 나쁜 것 보다 낫다.)

 

b. Learning Rate = 0.00005

 

C. Opimizer = Adam

 

D. Epochs = 100

학습 시에 과적합(오버피팅, overfitting)을 방지하기 위해 validation loss을 기준으로 조기 종료를 설정한다.

 

E. Performance Evaluation Indicator

성능 평가 지표는 Dice Coefficient Jaccard 사용하여, 각 폴딩 때마다 성능 평가한 것을 평균을 내서 전체적인 성능을 평가한다.

 

 

아래 class MODEL 이전글 "데이터셋 준비"와 "모델 설정"에서 설명한 코드를 모두 아우러서 사용하는 코드이다.

 

멤버 메소드인 Run_training은 K-Fold Cross Validation으로 모델을 학습하여 Training dataset에 대한 모델 평가를 보여준다.

Evaluation은 K-Fold Cross Validation으로 모델을 학습된 모델 중 최고 성능 모델을 선택하여 Test dataset에 대한 평가를 보여주고, Test dataset 일부 데이터에 대해 예측한 결과를 출력한다.

class MODEL():
    def __init__(self,
                 model_dir = MODEL_DIR,
                 batch_size = BATCH_SIZE,
                 width = WIDTH,
                 height = HEIGHT,
                 k = 0,
                ):
        
        self.batch_size = batch_size
        self.loss_fn = CategoricalFocalLoss()

        self.optimizer =  tf.keras.optimizers.Adam(learning_rate=0.00005)
        self.generator = Dataset_Generator()
        self.model_dir = MODEL_DIR
        self.optimal_k = k

        self.test_dataset = tf.data.Dataset.from_generator(
                        self.generator.test_generator,
                        (tf.float32, tf.int32),
                        (tf.TensorShape([1, HEIGHT, WIDTH, 3]), tf.TensorShape([1,  HEIGHT, WIDTH])),
                        )
        self.width = width
        self.height = height
    
    def __del__(self):
        print("MODEL is destructed")
    def Run_training(self, epochs= EPOCHS):        
        print("Mdoel Complie....")        
        
        # K-fold: k = 5
        K = 5
        mean_Dice = mean_IoU = 0
        DiceIoU_list = []
        for k in range(1, K+1):
            model = Unet()
            model.compile(loss = self.loss_fn, 
                          optimizer = self.optimizer,
                          metrics = [Dice, Jaccard]
                          )
            callbacks_list = [tf.keras.callbacks.ModelCheckpoint(
                                        filepath=os.path.join(
                                            f"{self.model_dir}U-Net_{k}.h5"),
                                        #monitor="val_loss",
                                        monitor="val_Dice",
                                        mode = "max",
                                        save_best_only=True,
                                        save_weights_only=True,
                                        verbose=1,
                                        ),
                              tf.keras.callbacks.EarlyStopping(
                                        #monitor = 'val_loss',
                                        monitor="val_Dice",
                                        mode = "max",
                                        min_delta = 0.01,
                                        patience = 5,
                                        )
                              ]
            print(f"{k}th fold Start Training....")
            
            history = model.fit(self.generator.train_generator(k),
                                steps_per_epoch = (K-1) * 20 * 8,
                                validation_data = self.generator.valid_generator(k),
                                validation_steps = 20 * 8,
                                callbacks = callbacks_list,
                                epochs = epochs,
                                batch_size = self.batch_size,
                                shuffle = True,
                                )
            
            loss = history.history['loss']
            val_loss = history.history['val_loss']
            dice = history.history["Dice"]
            val_dice = history.history["val_Dice"]
            iou = history.history["Jaccard"]
            val_iou = history.history["val_Jaccard"]
            
            DiceIoU_list.append( val_dice[-1] + val_iou[-1] )
            mean_Dice += val_dice[-1]
            mean_IoU += val_iou[-1]

            epochs_range = range(len(loss))
            
            plt.figure(k, figsize=(15, 5))

            plt.subplot(1, 3, 1)
            plt.plot(epochs_range, loss, label='Training Loss')
            plt.plot(epochs_range, val_loss, label='Validation Loss')
            plt.legend(loc='upper right')
            plt.title('Loss')

            plt.subplot(1, 3, 2)
            plt.plot(epochs_range, dice, label='Training Dice')
            plt.plot(epochs_range, val_dice, label='Validation Dice')
            plt.legend(loc='lower right')
            plt.title('Dice Coefficient')

            plt.subplot(1, 3, 3)
            plt.plot(epochs_range, iou, label='Training IoU')
            plt.plot(epochs_range, val_iou, label='Validation IoU')
            plt.legend(loc='lower right')
            plt.title('IoU')
            plt.show()

            input_image = tf.keras.Input(shape=(self.width, self.height, 3), name="image")
            predictions = model(input_image, training = True)
            inference_model = tf.keras.Model(inputs=input_image, outputs=predictions)

            for i, test in enumerate(self.test_dataset):
                img, mask = test
                prediction = inference_model.predict(img)
                
                img = img[0].numpy()
                mask = mask[0].numpy()

                prediction = prediction[0]
                prediction = tf.math.argmax(prediction, -1)
                prediction = prediction.numpy()
                
                fig = plt.figure(10, figsize = (20,20))
                ax1 = fig.add_subplot(1, 3, 1)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                ax1.imshow(img)
                ax1.set_title('Image')
                ax1.axis("off")

                ax2 = fig.add_subplot(1, 3, 2)
                ax2.imshow(mask)
                ax2.set_title('Ground Truth Mask')
                ax2.axis("off")

                ax3 = fig.add_subplot(1, 3, 3)
                ax3.imshow(prediction)
                ax3.set_title('Prediction')
                ax3.axis("off")
                plt.show()

                if i == 1:
                    break
            del model
        print("Training End\n\n")
        self.optimal_k = DiceIoU_list.index(max(DiceIoU_list)) + 1
        print(f"K-Fold Cross Validation Result\nmDice : {mean_Dice*20:.3f}, mIoU : {mean_IoU*20:.3f}, Optimal_K : {self.optimal_k}\n\n")
    def Evaluation(self, num_sample):
        input_image = tf.keras.Input(shape=(self.width, self.height, 3), name="image")
        model = Unet()
        model.load_weights(
            f"{self.model_dir}U-Net_{self.optimal_k}.h5")

        model.compile(loss = self.loss_fn, 
                      optimizer = self.optimizer,
                      metrics = [Dice, Jaccard]
                      )

        _, dice, iou = model.evaluate(self.test_dataset, batch_size = self.batch_size, verbose= 1)
        print(f"\n\nDice : {dice*100:.3f}, IoU : {iou*100:.3f}\n\n")
        predictions = model(input_image, training=True)
        inference_model = tf.keras.Model(inputs=input_image, outputs=predictions)
        
        print("Display predictions")
        for i, test in enumerate(self.test_dataset):
            img, mask = test
            prediction = inference_model.predict(img)
            
            img = img[0].numpy()
            mask = mask[0].numpy()

            prediction = prediction[0]
            prediction = tf.math.argmax(prediction, -1)
            prediction = prediction.numpy()
            
            fig = plt.figure(i, figsize = (20,20))
            ax1 = fig.add_subplot(1, 3, 1)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            ax1.imshow(img)
            ax1.set_title('Image')
            ax1.axis("off")

            ax2 = fig.add_subplot(1, 3, 2)
            ax2.imshow(mask)
            ax2.set_title('Ground Truth Mask')
            ax2.axis("off")

            ax3 = fig.add_subplot(1, 3, 3)
            ax3.imshow(prediction)
            ax3.set_title('Prediction')
            ax3.axis("off")
            plt.show()

            if i == num_sample:
                break

 

실행 코드는 아래와 같다.

# VGG16 + CFLoss
vgg_CFL = MODEL()
vgg_CFL.Run_training()
vgg_CFL.Evaluation(num_sample = -1)
del vgg_CFL

 

 

3.2. Result

학습은 각 폴드마다 15 ~ 25 epochs 에서 종료되었다.

훈련 데이터셋에서의 모델 성능 결과는 다음과 같다.

 

mDice = 90.444, mIoU : 84.016

 

이는 5번의 검증 결과를 평균낸 결과이다.

 

테스트 데이터셋에 대한 검증은 5-fold cross validation에서 가장 좋은 성능을 보인 두 번째 모델로 검증한다.

그 결과 테스트 데이터셋에 대한 모델 성능 결과는 다음과 같다.

 

Dice : 94.959, IoU : 90.995

 

아래 Figure 8은 테스트 데이터에 대한 원본 영상, Ground truth mask, 그리고 예측 결과이다.

 

테스트 데이터를 살펴본 결과 크기가 큰 exudates에 대한 예측은 생각보다 어렵다는 것을 알 수 있었을 뿐만 아니라, 본 게시글에서 사용한 데이터셋이 두 종류이기에 학습 시에 좋은 영향을 못 받은 느낌이 들기도 하였다.

 

또한 매우 작은 exudates에 대한 검출은 어려울 것으로 예측을 하고, 일반적인 dice loss function, jaccard loss function보다는 잘 학습이 되도록 focal loss function을 사용했는데, 역시나 작은 객체에 대한 검출은 매우 힘들었다.

 

무엇보다 데이터 부족은 본 게시글에서 근본적인 문제였기에, 이를 해결하는 방법은 multi-class semantic segmatation하기 보다는 각 클래스 별로 binary-class semantic segmantation을 진행하는 것이 더 좋은 방향이 될 것 같다.

 

또한, 병변을 segmentation으로 찾기 보다는, 분류(Classification)으로 CAM(Class Activation Map)을 통해서 병변을 찾는 방법도 알아볼 예정이다.

 

Figure 8. 테스트 데이터셋에 대한 결과