PSPNet의 학습 및 검증을 구현하고 실행해보자. 처음부터 PSPNet을 학습시키지 않고 학습된 모델을 사용하여 파인튜닝한다.
- 데이터 준비
학습된 모델로 파인튜닝을 실행한다. pspnet50_ADE20K.pth를 다운로드한다. 학습한 PSPNet의 결합 파라미터를 초깃값으로 하여 VOC2012 데이터셋을 파인튜닝한다.
- 학습 및 검증 구현
# 패키지 import
import random
import math
import time
import pandas as pd
import numpy as np
import torch
import torch.utils.data as data
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import torch.optim as optim# 초기설정
# Setup seeds
torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)
from utils.dataloader import make_datapath_list, DataTransform, VOCDataset
# 파일 경로 리스트 작성
rootpath = "./data/VOCdevkit/VOC2012/"
train_img_list, train_anno_list, val_img_list, val_anno_list = make_datapath_list(
rootpath=rootpath)
# Dataset 작성
# (RGB) 색의 평균값과 표준편차
color_mean = (0.485, 0.456, 0.406)
color_std = (0.229, 0.224, 0.225)
train_dataset = VOCDataset(train_img_list, train_anno_list, phase="train", transform=DataTransform(
input_size=475, color_mean=color_mean, color_std=color_std))
val_dataset = VOCDataset(val_img_list, val_anno_list, phase="val", transform=DataTransform(
input_size=475, color_mean=color_mean, color_std=color_std))
# DataLoader 작성
batch_size = 8
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}네트워크 모델을 만들기 위해 먼저 ADE20K 네트워크 모델을 준비한다. 출력 클래스 수는 ADE20K 데이터셋에 맞춰 150이다. 이 모델에 학습된 파라미터 pspnet50_ADE20k.pth를 읽는다. 최종 출력층을 Pascal VOC의 21 클래스로 하기 위하여 Decoder와 AuxLoss 모듈의 분류용 합성곱 층을 바꾼다. 21 클래스에 대응한 PSPNet이 된다.
from utils.pspnet import PSPNet
# 파인 튜닝으로 PSPNet을 작성
# ADE20K 데이터 세트의 학습된 모델을 사용하며, ADE20K는 클래스 수가 150입니다
net = PSPNet(n_classes=150)
# ADE20K 학습된 파라미터를 읽기
state_dict = torch.load("./weights/pspnet50_ADE20K.pth")
net.load_state_dict(state_dict)
# 분류용의 합성곱층을, 출력수 21으로 바꿈
n_classes = 21
net.decode_feature.classification = nn.Conv2d(
in_channels=512, out_channels=n_classes, kernel_size=1, stride=1, padding=0)
net.aux.classification = nn.Conv2d(
in_channels=256, out_channels=n_classes, kernel_size=1, stride=1, padding=0)
# 교체한 합성곱층을 초기화한다. 활성화 함수는 시그모이드 함수이므로 Xavier를 사용한다.
def weights_init(m):
if isinstance(m, nn.Conv2d):
nn.init.xavier_normal_(m.weight.data)
if m.bias is not None: # 바이어스 항이 있는 경우
nn.init.constant_(m.bias, 0.0)
net.decode_feature.classification.apply(weights_init)
net.aux.classification.apply(weights_init)
print('네트워크 설정 완료: 학습된 가중치를 로드했습니다')
다 클래스 분류의 손실함수인 크로스 엔트로피 오차 함수로 손실함수를 구현한다. 메인 손실과 AuxLoss 손실 합을 총 손실로 한다. AuxLoss는 계수 0.4를 곱하여 가중치를 메인 손실보다 작게 한다.
# 손실함수 정의
class PSPLoss(nn.Module):
"""PSPNet의 손실함수 클래스입니다"""
def __init__(self, aux_weight=0.4):
super(PSPLoss, self).__init__()
self.aux_weight = aux_weight # aux_loss의 가중치
def forward(self, outputs, targets):
"""
손실함수 계산
Parameters
----------
outputs : PSPNet의 출력(tuple)
(output=torch.Size([num_batch, 21, 475, 475]), output_aux=torch.Size([num_batch, 21, 475, 475]))。
targets : [num_batch, 475, 475]
정답 어노테이션 정보
Returns
-------
loss : 텐서
손실값
"""
loss = F.cross_entropy(outputs[0], targets, reduction='mean')
loss_aux = F.cross_entropy(outputs[1], targets, reduction='mean')
return loss+self.aux_weight*loss_aux
criterion = PSPLoss(aux_weight=0.4)
- 스케줄러로 에폭별 학습 비율 변경
마지막으로 파라미터 최적화 기법을 정의한다. 파인튜닝이기 때문에 입력에 가까운 모듈의 학습률은 작게, 교체한 합성곱 층을 가진 Decoder와 AuxLoss 모듈은 크게 설정한다.
이번에는 에폭에 따라 학습률을 변화시키는 스케줄러를 활용한다. 코드 scheduler= optim.lr_scheduler.LambdaLR(optimzier,lr_lambda = lambda_epoch)에 정의되어 있다. lambda_epoch 함수의 내용에 따라 옵티마이저 인스턴스의 학습률을 변화시키는 명령이다. lambda_epoch 함수는 최대 에폭 수를 30으로 하고 에폭을 거칠 때마다 학습률이 서서히 작아지도록 한다. return 하는 값을 옵티마이저 학습률에 곱한다. 스케줄러 학습률을 변화시키려면 네트워크 학습시 scheduler.step()을 실행해야 한다.
# 파인 튜닝이므로, 학습률은 작게
optimizer = optim.SGD([
{'params': net.feature_conv.parameters(), 'lr': 1e-3},
{'params': net.feature_res_1.parameters(), 'lr': 1e-3},
{'params': net.feature_res_2.parameters(), 'lr': 1e-3},
{'params': net.feature_dilated_res_1.parameters(), 'lr': 1e-3},
{'params': net.feature_dilated_res_2.parameters(), 'lr': 1e-3},
{'params': net.pyramid_pooling.parameters(), 'lr': 1e-3},
{'params': net.decode_feature.parameters(), 'lr': 1e-2},
{'params': net.aux.parameters(), 'lr': 1e-2},
], momentum=0.9, weight_decay=0.0001)
# 스케쥴러 설정
def lambda_epoch(epoch):
max_epoch = 30
return math.pow((1-epoch/max_epoch), 0.9)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_epoch)
# 모델을 학습시키는 함수를 작성
def train_model(net, dataloaders_dict, criterion, scheduler, optimizer, num_epochs):
# GPU가 사용 가능한지 확인
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("사용 장치: ", device)
# 네트워크를 GPU로
net.to(device)
# 네트워크가 어느 정도 고정되면 고속화한다
torch.backends.cudnn.benchmark = True
# 화상의 매수
num_train_imgs = len(dataloaders_dict["train"].dataset)
num_val_imgs = len(dataloaders_dict["val"].dataset)
batch_size = dataloaders_dict["train"].batch_size
# 반복자의 카운터 설정
iteration = 1
logs = []
# multiple minibatch
batch_multiplier = 3
# epoch 루프
for epoch in range(num_epochs):
# 시작 시간 저장
t_epoch_start = time.time()
t_iter_start = time.time()
epoch_train_loss = 0.0 # epoch의 손실합
epoch_val_loss = 0.0 # epoch의 손실합
print('-------------')
print('Epoch {}/{}'.format(epoch+1, num_epochs))
print('-------------')
# epoch별 훈련 및 검증 루프
for phase in ['train', 'val']:
if phase == 'train':
net.train() # 모델을 훈련 모드로
scheduler.step() # 최적화 scheduler 갱신
optimizer.zero_grad()
print('(train)')
else:
if((epoch+1) % 5 == 0):
net.eval() # 모델을 검증 모드로
print('-------------')
print('(val)')
else:
# 검증은 다섯 번 중에 한 번만 수행
continue
# 데이터 로더에서 minibatch씩 꺼내 루프
count = 0 # multiple minibatch
for imges, anno_class_imges in dataloaders_dict[phase]:
# 미니배치 크기가 1이면 배치 노멀라이제이션에서 오류가 발생하므로 회피
if imges.size()[0] == 1:
continue
# GPU가 사용가능하면 GPU에 데이터를 보낸다
imges = imges.to(device)
anno_class_imges = anno_class_imges.to(device)
# multiple minibatch로 파라미터 갱신
if (phase == 'train') and (count == 0):
optimizer.step()
optimizer.zero_grad()
count = batch_multiplier
# 순전파(forward) 계산
with torch.set_grad_enabled(phase == 'train'):
outputs = net(imges)
loss = criterion(
outputs, anno_class_imges.long()) / batch_multiplier
# 훈련시에는 역전파
if phase == 'train':
loss.backward() # 경사 계산
count -= 1 # multiple minibatch
if (iteration % 10 == 0): # 10iter에 한 번, loss를 표시
t_iter_finish = time.time()
duration = t_iter_finish - t_iter_start
print('반복 {} || Loss: {:.4f} || 10iter: {:.4f} sec.'.format(
iteration, loss.item()/batch_size*batch_multiplier, duration))
t_iter_start = time.time()
epoch_train_loss += loss.item() * batch_multiplier
iteration += 1
# 검증 시
else:
epoch_val_loss += loss.item() * batch_multiplier
# epoch의 phase별 loss와 정답률
t_epoch_finish = time.time()
print('-------------')
print('epoch {} || Epoch_TRAIN_Loss:{:.4f} ||Epoch_VAL_Loss:{:.4f}'.format(
epoch+1, epoch_train_loss/num_train_imgs, epoch_val_loss/num_val_imgs))
print('timer: {:.4f} sec.'.format(t_epoch_finish - t_epoch_start))
t_epoch_start = time.time()
# 로그 저장
log_epoch = {'epoch': epoch+1, 'train_loss': epoch_train_loss /
num_train_imgs, 'val_loss': epoch_val_loss/num_val_imgs}
logs.append(log_epoch)
df = pd.DataFrame(logs)
df.to_csv("log_output.csv")
# 최후의 네트워크를 저장
torch.save(net.state_dict(), 'weights/pspnet50_' +
str(epoch+1) + '.pth')
'Semantic Segmentation' 카테고리의 다른 글
| 8. 시맨틱 분할 추론 (0) | 2021.11.21 |
|---|---|
| 6. Decoder, AuxLoss 모듈 (0) | 2021.11.21 |
| 5. Pyramid Pooling 모듈 (0) | 2021.11.20 |
| 4. Feature 모듈 설명 및 구현(ResNet) (0) | 2021.11.20 |
| 3. PSPNet 네트워크 구성 및 구현 (0) | 2021.11.20 |