Pytorch 코드 작성 팁

 

Tensorflow 1.x 버전으로 코드를 작성하다가. Naver AI HackaThon을 계기로 Pytorch로 넘어오게 되었다. 어떻게 보면 비슷하지만 약간은 다른 Pytorch를 그동안 접하면서. 초보자들에게 유용한 팁을 한번 간단하게 정리를 해보고자 한다.

(이미 잘 하신 분들은 스킵하셔도 됩니다~~)

---

글의 구성은 다음과 같이 구성이 되어있습니다.

• Custom Dataset 을 Class로 만들어서 Data Loader 시키기.
• Transfer Learning and fine tunning with Pytorch
• Save and Load pickle file
• Model feature extractor

---

1. Custom Dataset Loader

Pytorch에서 기본적으로 제공해주는 Fashion MNIST, MNIST, Cifar-10 등. 이런 데이터셋은 코드 한줄로 딱 불러오면 손 쉽게 데이터를 불러올 수 있다. 하지만 실제로 딥러닝 관련 개발을 할때는 local에 있는 Data를 직접 불러와야한다. 이 때, pytorch 에서는 torchvision.dataset.ImageFolder, torchvision.dataset.DatasetFolder 과 같은 데이터를 불러오는 함수들이 존재한다. 하지만, 이런 함수를 적용하려면 폴더 구조도 해당 함수가 적용될 수 있는 폴더 구조로 되어있어야 한다는 단점이 발생된다. (물론 장점은 더 빠르고 손 쉽게 할 수 있다는 장점도 있다)

 

그래서 나는 보통 Dataset을 읽어오는 Class를 만들고, 이를 DataLoader를 통해서 Model에 넣기 위한 준비 과정을 한다.

 

1-1. 데이터 파일이 있는 경로를 넣어주는 경우.

1) Custom Data Class

class Custom_Dataset(Dataset):
    def __setup_files(self):
        files = glob(os.path.join(self.path, '*.jpg'))
        return files
    
    def __init__(self, path, res):
        self.path = path
        self.res = res
        self.files = self.__setup_files()
        
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, idx):
        print("="*20)
        img = cv2.imread(self.files[idx])
        im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(im_rgb, (self.res, self.res))
        
        return img

2) Dataset Class 호출

dataset = Custom_Dataset(path=path, res=img_size)
sample_img = next(iter(dataset))

print(type(sample_img)) # numpy.ndarry

3) DataLoader를 사용해서 torch.tensor 형태로 return

data_loader = DataLoader(dataset, batch_size=3, shuffle=True)
sample_loader = next(iter(data_loader))

print("sample_loader data type: ", type(sample_loader)) # <class 'torch.Tensor'>
print("sample_loader data size: ", sample_loader.size()) # torch.Size([3, 128, 128, 3])

1-2. 데이터 파일 경로가 담긴 list 형태로 넣어주는 경우.

1) Custom Data Class

class Custom_Dataset(Dataset):
    def __init__(self, file_list, res):
        self.file_list = file_list
        self.res = res
        
    def __len__(self):
        return len(self.file_list)
    
    def __getitem__(self, idx):
        print("="*20)
        img = cv2.imread(self.file_list[idx])
        im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(im_rgb, (self.res, self.res))
        
        return np.transpose(img, (2, 1, 0)) 
        # -> pytorch는 model에 들어갈 때 (N, C, W, H) 순서로 들어가기 때문에 여기서 transpose해줌.
        

2) 파일 경로를 list로 만들기 - 1

file_list = []
for (dir_path, dir_folder, file_names) in os.walk(path):
    for file_name in file_names:
        if file_name.split('.')[-1] == 'jpg':
            file_list.append(os.path.join(dir_path, file_name))

print("이미지 파일의 개수 : ", len(file_list))

3) 파일 경로를 list로 만들기 -2

path_list = glob(os.path.join(path, '*.jpg'))

print("이미지 파일의 개수 : ", len(file_list))

4) Dataset Class 호출

dataset = Custom_Dataset(file_list=file_list, res=img_size)

sample_img = next(iter(dataset))
type(sample_img)

5) DataLoader

data_loader = DataLoader(dataset, batch_size=3, shuffle=True)
sample_loader = next(iter(data_loader))

print("sample_loader data type: ", type(sample_loader)) # <class 'torch.Tensor'>
print("sample_loader data size: ", sample_loader.size()) # torch.Size([3, 3, 128, 128])

1-3. torchvision.transforms을 사용한 경우

1) Custom Data Class

class Custom_Dataset(Dataset):
    def __init__(self, file_list, transforms=None):
        self.file_list = file_list
        self.transforms = transforms
        
    def __len__(self):
        return len(self.file_list)
    
    def __getitem__(self, idx):
        img = cv2.imread(self.file_list[idx])
        im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(im_rgb.transpose(1, 0, 2))
        
        if self.transforms is not None:
            img = self.transforms(img)
        
        return img
        

2) Dataset 호출 및 DataLoader

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(), # (H, W, C) --> (C, H, W)로 변경해줌.
])

dataset = Custom_Dataset(file_list=file_list, transforms=transform)
data_loader = DataLoader(dataset, batch_size=3, shuffle=True)

---

2. Transfer Learning

Transfer Learning(전이학습)은 미리 누군가가 엄청나게 많은 Dataset(보통 ImagNet)으로 오랜시간 공들여서 학습한 모델을 그대로 가지고오는 방법이다. Tensorflow, Keras와 동일하게 Pytorch에서도 동일하게 제공해준 모델을 가지고와서 이어서 사용할 수 있다.

사용할 수 있는 방법은 크게 2가지 방법이다. 

• ....pretrained=True) --> 미리 학습된 weight들을 이어서 가지고온다.
• ....pretrained=False) --> 미리 학습된 weight들은 안가지고 오고, random하게 weight를 부여한 후 model 구조만 가지고 온다

2-1. Model 모델 불러오기

import torchvision.models as models

""" pytorch 에서 사용 가능한 model """
# resnet18 = models.resnet18(pretrained=True)
# alexnet = models.alexnet(pretrained=True)
# squeezenet = models.squeezenet1_0(pretrained=True)
# vgg16 = models.vgg16(pretrained=True)
# densenet = models.densenet161(pretrained=True)
# inception = models.inception_v3(pretrained=True)
# googlenet = models.googlenet(pretrained=True)
# shufflenet = models.shufflenet_v2_x1_0(pretrained=True)
# mobilenet = models.mobilenet_v2(pretrained=True)
# resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
# wide_resnet50_2 = models.wide_resnet50_2(pretrained=True)
# mnasnet = models.mnasnet1_0(pretrained=True)

model = models.inception_v3(pretrained=True)
print(model)

모델을 불러오면, Classifier 부분(보통 Fully connected layer부분)을 논문 상 그대로 가지고 온다. 이 부분은 내가 가지고 있는 Data들로 해결해야한 문제마다 output이 다 다르다. 따라서 이 부분을 설정해주어야 하는데 이를 Fine-tunning (미세조정) 이라고 한다. 한번 Fully connected layer 마지막 output node 부분을 변경 해보도록 하자.

Classifier 부분같은 경우. 예전에 나온 모델들을 거의 Fully connected layer를 많이 사용했는데 요즘 나온 논문들을 보면 Fully connected layer 부분을 사용하면 Parameter 수가 많아지다 보니 거의 안쓰는 모델들을 많이 볼 수 있다. 그래서 Fully connect layer 부분은 제거해주고 그 전 conv layer 부분까지만 이어 불러와서 사용하는 경우가 많다.

2-2. Fine tunning(마지막layer 부분만.)

1) Model의 Fully-connected layer의 input features 부분 뽑아내기.

FC_in_features = model.fc.in_features

print("number of Fully connected layer input feature: ", FC_in_features)
# >> number of Fully connected layer input feature:  2048

2) Classifier 출력 node 변경

model.fc = nn.Linear(in_features=FC_in_features, out_features=2)

print(model.fc)
# >> Linear(in_features=2048, out_features=2, bias=True)

3) Model을 학습 가능한 상태로 만들기 위해서 CPU or GPU에 올리기

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("device: ", device)

model.to(device)

---

3. Save and load Pickle file

Pickle file은 우리가 dataset을 전처리 하고 뭐하고 뭐하고 하는 과정들을 매번 코드를 실행할때마다 이런 과정을 거친다면 메모리 소비도 많이 일어나고 그만큼 시간도 늘어나게 된다. 그래서 여기서는 Pickle 로 데이터를 저장하고 불러오는 방법을 소개해보고자 한다.

(물론 Pickle말고 .npy, .npz 등 여러가지 방법들이 존재하기는 한다.)

 

Pickle의 장점은, 데이터 처리를 해준 상태에서 그 완료된 상태를 Pickle로 저장하면 그 상태를 Binary로 저장을 한다. 그러면 쉽게 해당 데이터를 불러오기만 하면 전처리 과정 없이 바로 사용할 수 있다는 장점이 있다. 그러면 바로 코드를 확인해보자.

 

텍스트 상태의 데이터가 아닌 파이썬 객체 자체를 파일로 저장하는 것 입니다.

 

3-1. Numpy image를 Pickle로 만들기

import cv2
import os
import pickle
import gzip

img = cv2.imread(file_list[0])

with gzip.open('./sample_pickle.pickle', 'wb') as f:
    pickle.dump(img, f)

3-2. Pickle 불러오기

with gzip.open('./sample_pickle.pickle', 'rb') as f:
    img = pickle.load(f)
    print(img.shape)
    plt.imshow(img)

여기서 gzip을 사용하는 이유는, 압축을 해주는 형식이라서 데이터 용량을 더 줄여주면서 저장할 수 있다.

---

4. Model Feature Extraction

딥러닝 모델을 짜다가 가끔 이런 생각이 들 때가 있을것이다. "A라는 모델에서 5번째 레이어인 부분을 추출해서 어딘가에다가 넣고싶다. 근데 나는 pretrained된 모델을 사용해서 모델 구조를 내 마음대로 바꿀 수 없을거 같다." 라는 생각이 들때가 있을것이다. 그럴 때 사용하는 방법이다. 위에서 Fine tunning을 할 때 사용한 방법인데. 여기서도 모델 layer를 내 마음대로 접근해서 빼오고 추가하고 하는 방법을 설명하도록 하겠다.

1) Model 불러오기

import torchvision.models as models

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

model = models.vgg16()
model = model.to(device)

#model.load_state_dict(torch.load("dddd.pth"))

2) 중간 부분 layer의 Feature를 뽑아내기

new_feature_map = nn.Sequential(*list(model.features.children())[:-5]) 
print(new_feature_map)

3) Model에 새로 적용하기

model.features = new_feature_map
print(model)

 

 

해당 튜토리얼은 Github에 정리해서 올려보도록 하겠습니다.

 

 

 

Reference

- https://tutorials.pytorch.kr/beginner/transfer_learning_tutorial.html
- https://pytorch.org/docs/stable/torchvision/models.html