[Pytorch]pytorch를 활용한 EfficientNet 학습 코드 [3/3]

데이터 세트 구성하는 코드

https://twobeach.tistory.com/61

[Pytorch]pytorch를 활용한 EfficientNet 학습 코드 [1/3]

학습 코드

https://twobeach.tistory.com/62

[Pytorch]pytorch를 활용한 EfficientNet 학습 코드 [2/3]

이제 학습해서 나온 가중치를 이용한 테스트를 진행하는 코드를 작성해봅시다.

학습해서나온 EfficientNet best model pth 파일을 가지고와서 테스트 데이터 세트의 accuracy, error rate, f1-score를 뽑아 봅시다.

 

먼저 테스트 위의 링크를 참고하여 데이터 세트를 학습 데이터 경로를 저장하게 만드는 코드를 통해 똑같은 과정을 거쳐 테스트 데이터 세트 경로를 저장하는 csv파일을 만들어줍시다.

 

그 후 이제 사용할 라이브러리를 import 시켜줍니다

 

import os ,time

import glob
import random
from PIL import Image, ImageFile
import ttach as tta
import matplotlib.pyplot as plt
import imutils

import torch
import torch.nn as nn
from torchvision import transforms
import torchvision.transforms.functional as TF
from torch.utils.data import Dataset, DataLoader
from torch.cuda.amp import autocast, grad_scaler
from torch.optim.lr_scheduler import _LRScheduler
import torch_optimizer as optim
from efficientnet_pytorch import EfficientNet

import numpy as np
import pandas as pd, cv2
from tqdm import tqdm
from easydict import EasyDict
from sklearn.preprocessing import LabelEncoder #Label Encoder로 숫자로 변경함
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import f1_score
from sklearn.metrics import accuracy_score
from ptflops import get_model_complexity_info

from collections import Counter
import argparse
import neptune.new as neptune
import timm

#안하면 오류남
import warnings
from pandas.core.common import SettingWithCopyWarning
warnings.simplefilter(action="ignore", category=SettingWithCopyWarning)
warnings.filterwarnings(action='ignore')

 

정상적으로 라이브러리가 import 되었다면 이제 cuda 사용여부를 결정해봅니다. 여기서 cuda device를 사용하지 못하고 cpu로 뜨는경우 cuda를 Nvidia 홈페이지를 통해 설치하시길 바랍니다.

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"device setting : {device}")

 

이제 만들어진 테스트 데이터 세트 경로 csv 파일을 변수에 저장합니다.

test_df = pd.read_csv('/home/user/effcientNet/explosive/RP_explosive_test/test_df.csv')

 

그 후 사용할 EfficientNet 사전학습 모델 구조를 불러옵니다.

 

args = EasyDict({'encoder_name':'efficientnet-b0',
                 'drop_path_rate':0.2,
                 'num_classes':17, # class는 본인의 class 개수에 맞추시길 바랍니다.
                 'bad' : False
                })

이제 테스트 데이터 세트 이미지를 형식에 맞춰 집어넣는 코드를 작성해봅시다.

논문에서 제공하는 값들로 구성이 되었기때문에 이 함수에 있는 값들은 변경하시는 것을 추천드리지 않습니다!

def get_train_augmentation(img_size, ver):
    if ver==1:
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((img_size, img_size)),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
                ])      
    return transform

transform = get_train_augmentation(img_size = 224, ver = 1)

이제 기본적인 테스트를 위한 함수들을 작성해줍니다.

class Network(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.encoder = EfficientNet.from_pretrained(args.encoder_name,num_classes = args.num_classes )
       
    def forward(self, x):
        x = self.encoder(x)
        return x

class Test_Dataset(Dataset):
    def __init__(self, df, transform=None):
        self.file_name = df['file_name'].values
        self.transform = transform

        print(f"테스트 데이터셋 생성 완료,,\n 데이터셋 사이즈 : {len(self.file_name)}")

    def __getitem__(self, index):        
        image = np.array(Image.open(f'{self.file_name[index]}').convert('RGB'))

        if self.transform is not None:
            image = self.transform(Image.fromarray(image))

        return image

    def __len__(self):
        return len(self.file_name)


test_dataset = Test_Dataset(test_df, transform)
test_load = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=4)

test_load 부분은 본인이 학습할 때 사용하였던 batch_size로 변경해줍니다 저의 경우는 64를 사용하였습니다.

 

이제 테스트를 진행해 봅시다!

train이 끝나고 나온 모델 가중치를 넣어주고 정확도와 F1-score를 산출하는 코드입니다.

def predict(args, test_load, model_path):
    _transforms = tta.Compose([
        tta.Rotate90(angles=[0,20]),
    ])
    model = Network(args).to(device)
    model.load_state_dict(torch.load(model_path)['state_dict'])
    model = tta.ClassificationTTAWrapper(model, _transforms).to(device)
    model.eval()

    output = []
    pred = []
    with torch.no_grad():
        with autocast():
            for batch in tqdm(test_load):
                images = torch.tensor(batch, dtype = torch.float32, device = device).clone().detach()
                preds = model(images)
                pred.extend(preds)
                output.extend(torch.tensor(torch.argmax(preds, dim=1), dtype=torch.int32).cpu().numpy())
   
    return  output, pred

model_path = "./result/001/efficientnet-b0_best_model.pth"
args = EasyDict({'encoder_name':'efficientnet-b0',
                 'drop_path_rate':0.2,
                 'num_classes':17,
                })
output1, preds1 = predict(args, test_load, model_path)

answer_labels = test_df['label']
submit_labels = output1

print(f1_score(answer_labels,submit_labels,average='macro'))
print(accuracy_score(answer_labels,submit_labels))
print(len(answer_labels))

 

이제 테스트가 끝났으나 하나더 코딩을 진행해보자면 테스트 진행중 틀렸던 파일의 경로가 무엇이며 label을 뭐로 잘못 인식했는지 확인해보는 코드입니다.

for i in range(len(answer_labels)):
   
    if answer_labels[i] != submit_labels[i]:
        print(f'정답 : {answer_labels[i]}-> 예측 : {submit_labels[i]}')
        print(f"Wrong Answer : {test_df['file_name'][i]}")
        image = cv2.imread(test_df['file_name'][i])
        image = imutils.resize(image, width=800)
        cv2.imshow(f'correct : {answer_labels[i]}-> predict : {submit_labels[i]}',image)
        cv2.waitKey()
        cv2.destroyAllWindows()

 

이로써 EfficientNet의 데이터 세트 구성하는 것과 학습법부터 테스를 진행하는 것까지 해보았습니다.

전체적인 리뷰로는 b7이 b0보다 모델의 크기가 크다고해서 꼭 좋은 결과를 얻지는 못한 것처럼 본인이 가진 데이터 세트에 맞는 사전학습 모델을 찾아야 할것이고 저가 진행했던 연구에서는 b1이 가장 좋은 결과를 얻었습니다.

 

모델 자체가 다른 모델에 비해 가볍지만 다른 모델과 견주어 봣을때도 정확도면에서 밀리지 않는 모습을 결과를 통해 알 수 있었습니다.

다들 이 코드를 통해 성공적인 연구 결과를 얻으시길 바라겠습니다.