pytorch 재활훈련 -fine tuning 구현해보기-

1. fine tuning

 

출력층 등을 변경한 모델을 학습된 모델을 기반으로 구축한 후, 직접 준비한 데이터로 신경망 모델의 결합 가중치를 학습시키는 방법

 

결합 가중치의 초기값은 학습된 모델의 parameter를 사용하는 것이 보통이다.

 

finetuning은 transfer learning과는 다르게 출력층, 출력층에 가까운 부분 뿐만 아니라, 모든 층의 parameter를 재학습시킨다는 것이 특징이다.

 

일반적으로 입력층에 가까운 부분의 parameter는 learning rate를 작게 설정하고, (경우에 따라서는 바꾸지 않고)

 

출력층에 가까운 부분의 parameter는 learning rate를 크게 설정한다

 

transfer learning처럼 학습된 모델을 기반으로 하는 fine tuning은 직접 준비한 데이터가 적어도 높은 성능의 딥러닝을 실현하기 쉽다는 장점이 있다.

 

 

2. 패키지 import

 

#패키지 import

import glob

import os.path as osp

import numpy as np
import random
from PIL import Image

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data

from torchvision import models,transforms

from tqdm import tqdm

#난수 설정
torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)

 

3. dataset, dataloader 작성

 

이전에 작성했던 클래스들을 이용해서 dataset과 dataloader 준비

 

이미지를 전처리하는 클래스 ImageTransform

 

#입력이미지의 전처리 클래스
#훈련시와 추론시 처리가 다르다

class ImageTransform():
    
    """
    이미지 전처리 클래스, 훈련시, 검증시 동작이 다르다
    이미지 크기를 리사이즈하고, 색상을 표준화

    훈련시에는 RandomResizedCrop과 RandomHorizontalFlip으로 데이터 확장

    attributes
    ----------------
    resize: int
        리사이즈 대상 이미지의 크기
    
    mean: (R,G,B)
        각 색상 채널의 평균값
    
    std: (R,G,B)
        각 색상 채널의 표준편차
    
    """

    def __init__(self, resize, mean, std):
        
        self.data_transform = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(
                    resize,scale = (0.5,1.0) #데이터 확장
                ),
                transforms.RandomHorizontalFlip(), #데이터 확장
                transforms.ToTensor(), #텐서 변환
                transforms.Normalize(mean,std) #표준화
            ]),
            'val': transforms.Compose([
                transforms.Resize(resize), #리사이즈
                transforms.CenterCrop(resize), #이미지 중앙을 resize
                transforms.ToTensor(), #텐서로 변환
                transforms.Normalize(mean,std) #표준화
            ])
        }
    
    def __call__(self, img, phase='train'):
        
        """
        Parameters
        ------------
        phase: 'train' or 'val'
            전처리 모드를 지정
        """

        return self.data_transform[phase](img)

 

데이터 경로 리스트를 담는 make_datapath_list

 

#개미와 벌이 담긴 이미지 파일 경로 리스트 작성하기
def make_datapath_list(phase='train'):
    
    """
    데이터의 경로를 저장한 리스트를 작성한다.

    parameters
    ----------------
    phase: 'train' or 'val'
        훈련 데이터 또는 검증 데이터를 지정
    
    Returns
    ---------------
    path_list: list
        데이터 경로를 저장한 리스트
    
    """

    rootpath = './data/hymenoptera_data/'
    target_path = osp.join(rootpath+phase+'/**/*.jpg')
    print(target_path)

    path_list = []#여기에 저장

    #glob을 이용하여 하위 디렉토리의 파일 경로 가져오기

    for path in glob.glob(target_path):
        
        path_list.append(path)
    
    return path_list

 

dataset을 구하는 클래스

 

#개미와 벌 이미지에 대한 dataset

class HymenopteraDataset(data.Dataset):
    
    """
    개미와 벌 이미지의 dataset class
    Pytorch의 dataset 클래스를 상속

    attributes
    -------------
    file_list: 리스트
        이미지 경로를 저장한 리스트
    
    transform: object
        전처리 클래스의 인스턴스
    
    phase: 'train' or 'test'
        학습인지 훈련인지를 결정
    
    """

    def __init__(self, file_list, transform=None, phase='train'):
        
        self.file_list = file_list #파일 경로 리스트
        self.transform = transform #전처리 클래스의 인스턴스
        self.phase = phase #train or val 지정
    
    def __len__(self):
        
        """이미지 개수를 반환"""
        return len(self.file_list)
    
    def __getitem__(self,index):
        
        """
        전처리한 이미지의 tensor 형식 데이터, 라벨 취득
        """

        #index번째 이미지 가져오기
        img_path = self.file_list[index]
        img = Image.open(img_path) #[높이][너비][색]

        #이미지 전처리
        img_transformed = self.transform(
            img, self.phase) #torch.Size([3,244,244])
        
        #이미지 라벨을 파일 이름에서 추출
        if self.phase == 'train':
            
            label = img_path[30:34]
        
        elif self.phase == 'val':
            
            label = img_path[28:32]
        
        #라벨을 숫자로
        if label == 'ants':
            
            label = 0
        
        elif label == 'bees':
            
            label = 1
        
        return img_transformed, label

 

dataset 인스턴스와 dataloader 작성하기

 

#개미와 벌 이미지 파일 경로 리스트
train_list = make_datapath_list(phase='train')
val_list = make_datapath_list(phase='val')

#dataset
size = 224
mean = (0.485,0.456,0.406)
std = (0.229,0.224,0.225)

train_dataset = HymenopteraDataset(
    file_list = train_list, transform = ImageTransform(size,mean,std), phase='train'
)
val_dataset = HymenopteraDataset(
    file_list = val_list, transform = ImageTransform(size,mean,std), phase='val'
)

#dataloader

batch_size = 32

train_dataloader = data.DataLoader(
    train_dataset,batch_size=batch_size, shuffle = True
)
val_dataloader = data.DataLoader(
    val_dataset, batch_size = batch_size, shuffle = False
)

#사전 객체에 정리
dataloaders_dict = {'train':train_dataloader, 'val':val_dataloader}

 

 

4. 모델 작성

 

학습된 vgg-16을 불러오고, 출력층을 개미,벌 2클래스로 바꾼다

 

#학습된 vgg-16 모델 로드

#vgg-16모델 인스턴스 생성
use_pretrained=True #사전학습된 모델 사용
net = models.vgg16(pretrained=use_pretrained)

#vgg16 마지막 출력층의 출력 유닛을 개미,벌 2개로
net.classifier[6] = nn.Linear(in_features=4096, out_features=2)

#훈련 모드로 설정하기
net.train()

print('네트워크 설정 완료: 학습된 가중치를 로드하고 훈련 모드로 설정했다')

 

5. 손실함수 정의

 

분류 문제이므로 cross entropy 함수를 사용

 

criterion = nn.CrossEntropyLoss()

 

 

6. 최적화 방법 설정하기

 

fine tuning은 모든 층의 parameter를 학습할 수 있도록 설정

 

각 층의 learning rate를 바꿀 수 있도록 parameter를 설정

 

vgg-16의 전반부 features 모듈의 parameter는 update_param_names_1 변수에

 

후반부 classifier module은 처음 2개의 fully connected layer는 update_param_names_2에

 

교체한 마지막 fully connected layer는 update_param_names_3에

 

그래서 각각 다른 learning rate를 적용하도록 만든다

 

#fine tuning으로 학습할 가중치를 params_to_update 변수의 1~3에 저장

params_to_update_1 = []
params_to_update_2 = []
params_to_update_3 = []

#학습시킬 층의 파라미터 명
update_param_names_1 = ['features']
update_param_names_2 = ['classifier.0.weight', 'classifier.0.bias','classifier.3.weight','classifier.3.bias']
update_param_names_3 = ['classifier.6.weight','classifier.6.bias']


#파라미터를 각 리스트에 저장

for name,param in net.named_parameters():
    
    if update_param_names_1[0] in name:
        
        param.requires_grad = True
        params_to_update_1.append(param)
        print("params_to_update_1에 저장: ",name)

    elif name in update_param_names_2:
        
        param.requires_grad = True
        params_to_update_2.append(param)
        print("params_to_update_2에 저장: ",name)
    
    elif name in update_param_names_3:
        
        param.requires_grad = True
        params_to_update_3.append(param)
        print("params_to_update_3에 저장: ",name)
    
    else:
        
        param.requires_grad = False
        print("gradient 계산없음. 학습하지 않음: ",name)

 

optimizer에 리스트로 parameter를 넣으면 각 층마다 다른 learning rate를 적용시킬 수 있다

 

#층마다 다른 learning rate를 적용시키는 최적화 기법

optimizer = optim.SGD([
    {'params': params_to_update_1, 'lr':1e-4},
    {'params': params_to_update_2, 'lr':5e-4},
    {'params': params_to_update_3, 'lr':1e-3},
],momentum=0.9)

 

momentum은 모든 parameter에 동일하게 적용된다

 

아마 해보진 않았지만 learning rate처럼 parameter dict에 momentum 다르게 적용할 수 있을듯

 

 

7. 학습 및 검증하기

 

이전에 사용했던 train_model 함수를 작성하고, gpu를 사용할 수 있도록 설정 추가

 

gpu를 사용할려면, device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')를 수행

 

gpu를 사용할 수 있는 경우는 device변수에 cuda:0가 저장되고 그렇지 않으면 cpu가 저장

 

device 변수를 사용해서 모델, 데이터, 라벨을 gpu에 올려둔다

 

.to(device)를 이용하면 gpu에 올려둘 수 있다

 

#모델을 학습시키는 함수 작성
def train_model(net, dataloaders_dict, criterion, optimizer, num_epochs):
    
    #초기 설정
    #gpu 사용가능한지 확인
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('사용 장치: ',device)

    #네트워크를 gpu에
    net.to(device)

    #네트워크가 어느 정도 고정되면, forward, loss에 대한 gpu 계산을 가속화시킴
    torch.backends.cudnn.benchmark = True

    #훈련
    for epoch in range(num_epochs):
        
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('---------------')

        #epoch별 훈련, 검증
        for phase in ['train','val']:
            
            if phase == 'train':
                
                net.train() #모델을 훈련모드로
            
            else:
                
                net.eval() #모델을 검증모드로
            
            epoch_loss = 0.0 #epoch 손실 합
            epoch_corrects = 0 #epoch 정답 수

            #미학습시의 검증 성능을 확인하기 위해 epoch=0의 훈련 생략
            if (epoch == 0 and phase == 'train'):
                
                continue
            
            #dataloader에서 미니 배치를 꺼내 반복
            for inputs, labels in tqdm(dataloaders_dict[phase]):
                
                #gpu가 사용가능하다면.. gpu에 데이터 보내기
                inputs = inputs.to(device)
                labels = labels.to(device)

                #optimizer를 초기화
                optimizer.zero_grad()

                #forward계산

                with torch.set_grad_enabled(phase=='train'):
                    
                    outputs = net(inputs)
                    loss = criterion(outputs, labels)#손실 계산
                    _,preds = torch.max(outputs,1) #라벨 예측

                    #훈련시에 역전파
                    if phase == 'train':
                        
                        loss.backward()
                        optimizer.step()
                    
                    #결과 계산
                    epoch_loss += loss.item()*inputs.size(0) #loss 합계 갱신

                    #정답 수의 합계 갱신
                    epoch_corrects += torch.sum(preds == labels.data)
            
            #epoch별 loss와 정답률
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double()/len(dataloaders_dict[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase,epoch_loss,epoch_acc
            ))

 

여기서 torch.backends.cudnn.benchmark = True는 신경망의 forward, loss 계산이 반복할때, 어느 정도 일정하다면 gpu 계산을 고속화시키는 설정이라고 함

 

실제 학습, 검증을 수행하면

 

#학습 및 검증
num_epochs = 2
train_model(net, dataloaders_dict, criterion, optimizer, num_epochs = num_epochs)

 

 

 

 

첫 epoch에서는 학습을 하지않고 continue로 건너 뛰었는데 그 결과 검증 정답률이 약 44%

 

처음에는 개미와 벌의 이미지를 능숙하게 분류하지 못한다

 

1epoch 학습하면, 학습 데이터 정답률이 71%, 검증 데이터 정답률이 96%가 된다

 

cuda:0로 gpu를 사용하여, 전체 계산이 약 21초만에 끝나 cpu보다 훨씬 빠르다

 

 

8. 학습된 모델 저장, 로드

 

학습한 모델을 저장하는 방법과 로드하는 방법

 

저장할 경우에는 모델을 저장한 net 변수를 .state_dict()를 활용해 parameter를 사전형으로 꺼내고

 

torch.save()로 저장

 

#pytorch model parameter 저장

save_path = './weights_fine_tuning.pth'
torch.save(net.state_dict(), save_path)

 

parameter는 보통 pth파일로 저장함

 

로드할 경우에는 torch.load()로 사전형 객체를 로드하여 네트워크를 load_state_dict()로 저장

 

모델을 로드할때는 당연하지만 로드할려는 모델과 동일한 구성의 네트워크 모델이어야함

 

vgg16이면 당연히 vgg16을 가져와야지 resnet을 가져오면 안된다는 소리

 

#pytorch network parameter load

load_path = './weights_fine_tuning.pth'
load_weights = torch.load(load_path)

#net은 로드할려는 모델과 동일한 구성의 모델이어야함
net.load_state_dict(load_weights)

 

gpu상에 저장한 파일을 cpu에 로드하고 싶으면 map_location을 사용해야한다고 함

 

#만약 gpu상에 저장된 가중치를 cpu에 로드하고 싶다면
load_weights = torch.load(load_path, map_location = {'cuda:0':'cpu'})
net.load_state_dict(load_weights)