[Kaggle] CT Medical Image - (3) 모델 만들기

CT Medical Image Classification

(0) 서론, 액션플랜
(1) DICOM 파일 array로 전환 및 시각화
(2) DataLoader 만들기

 

0. 필요한 라이브러리 임포트 및 데이터 로더

import pandas as pd
import numpy as np
import random

import torch, torchvision
import torchvision.models as models
import torch.nn as nn
from torch.utils.data import DataLoader

from customized_dataloader import CustomDataset,customsampler

df = pd.read_csv("final_df.csv")

df.loc[df['Contrast']==True,['Contrast']] = 1
df.loc[df['Contrast']==False,['Contrast']] = 0

df['Contrast'] = df['Contrast'].astype('int')

train_images = df['dicom_name'].values
train_labels = df['Contrast'].values

 

그리고 나중에 검증하기 위해서 train, test 데이터 셋 split하기

 

X_train, X_test, y_train, y_test = train_test_split(train_images, train_labels, test_size=0.33, random_state=42)

 

데이터로더

 

dataset = CustomDataset(X_train, y_train)
sampler = customsampler(y_train, 4,4)
dataloader = DataLoader(dataset, batch_size=8,sampler=sampler)

total_batch = len(dataloader)

 

1. 사전 훈련된 모델 다운로드 및 커스터마이징

 

device = 'cuda' if torch.cuda.is_available() else 'cpu'

vgg16 = models.vgg16(pretrained=True)

vgg16.features[0] = nn.Conv2d(1,64, kernel_size=(3,3))

vgg16.classifier[3] = nn.Linear(4096,512)
vgg16.classifier[6] = nn.Linear(512,1)

model = nn.Sequential(
        vgg16,
        nn.Sigmoid()
        )

model.to(device)

 

- DICOM 파일은 기본적으로 흑백이기에, input의 채널이 1뿐이다. 따라서 제일 처음 오는 합성곱 필터의 채널수를 3에서 1로 변경해준다. 

 

- 또한 FC layer 끝의 노드가 하나만 필요하기 때문에 전 Linear 레이어에서 4096개의 노드는 너무 많은 것 같아서 차츰 줄어들도록 했다. 

 

- BCE Loss를 사용했는데(추후에 언급하겠지만) sigmoid함수가 없으니 에러가 계속 발생했다. 

 sigmoid 함수를 통해서 latent vector의 범위를 0~1로 변경해주어야 한다. 

 

- 또한 기존의 네트워크 마지막에 추가로 layer를 추가하기 위해서 처음에는

vgg16.classifier[7] = nn.Sigmoid()

로 하였으나 에러가 발생했다. 리스트로 추가하는 것과 완전 똑같게는 되지 않는 모양이다. 그래서 위와 같이 기존의 model에 nn.Sequentail 함수를 이용해서 추가해야 한다. 

 

 2. 모델을 이용한 훈련

 

model.train()
criterion = nn.BCELoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

epoch = 10
print("before initial")
print(list(model.parameters())[0][0])

for i in range(epoch):
    avg_cost = 0.0
    for e, data in enumerate(dataloader):

        image, label = data
        image = image.reshape(16,512,512,1) #채널수가 흑백=>마지막에 1
        
        image = image.permute(0,3,1,2) #batch, channel, width, height 순으로 텐서 변경
        
        image = image.to(device)
        label = label.to(device)
        
        image.requires_grad = True
        label.requires_grad = False 
        
        model.train()
        optimizer.zero_grad()

        outputs = model(image)

        loss = criterion(outputs, label)
        loss.backward()
        optimizer.step()
        
        avg_cost += loss / total_batch
        
        print("after train")
        print(list(model.parameters())[0][0])
        if e == 1: break
    print(f'[Epoch:{i + 1}] cost = {avg_cost}')
print('Learning Finished!')

 

- 학습함에도 불구하고, loss가 줄어들지 않자 나는 멘붕이 왔다. 내가 겪었던 에러와 해결 방안을 추후를 위해 정리해 놓는다. 

 

(1) optimizer의 인자로 model.parameters()가 아니라 vgg16.parameters()를 넣어놓고 몰랐던 것 

(2) image의 requires_grad를 False로 해놓았던 것 

(3) Loss 역시 device로 보내놓지 않은 것

 

- 하지만 무엇보다도 중요한 건 위 문제를 발견하고 해결한 후, 왜 업데이트가 안 되는지 알기 위해서 모델의 파라미터가 바뀌지 않는지 일부 확인해보았는데,

알고보니까 parameter 자체는 잘 바뀌는데 loss가 안 줄어드는 것이 문제였다. 

 

머쓱.

 

하지만 성능이 안 좋아도 일단 baseline은 완성했으니 넘어가도록 하자.

agumentation까지 모두 다 해보는 것이 목표니까!

 

3. 모델 저장

 

torch.save(model.state_dict(), '/content/drive/MyDrive/video/archive/baseline.pth')