[파이토치로 시작하는 딥러닝 기초]10.3 ImageFolder / 모델 저장 / 모델 불러오기

ImageFolder

 

나만의 데이터 셋 준비하기

ImageFolder란? 

로컬에 저장된 이미지 데이터를 불러올 때 사용하는 pytorch 라이브러리

 

데이터를 준비할 때에는 명확하게 구분되는 사진을 사용해야 한다. 

위와 같이 구분하는 label의 class 개수에 따라 folder를 생성하고 그 안에 해당 라벨에 맞는 이미지를 삽입한다. 

 

데이터 불러오기 실습

import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader

from matplotlib.pyplot import imshow
%matplotlib inline

train_data = torchvision.datasets.ImageFolder(root = "data/custom_data/origin_data",
                                             transform = None)
                                             
for num, value in enumerate(train_data):
    data, label = value
    print(num, data, label)
    
    imshow(data)
    break                                             

사이즈 낮추기

512X256 사이즈는 너무 크기 때문에 이미지 사이즈를 낮춰야 한다.

trans = transforms.Compose([
    transforms.Resize([64, 128])
])

train_data = torchvision.datasets.ImageFolder(root = 'data/custom_data/origin_data',
                                             transform = trans)

for num, value in enumerate(train_data):
    data, label = value
    print(num, data, label)
    
    imshow(data)
    break

사진의 길이가 4분의 1로 줄어든 것을 확인할 수 있다. 

줄어든 사진을 별도 폴더에 저장한다. 

for num, value in enumerate(train_data):
    data, label = value
    print(num, data, label)
    
    if(label == 0):
        data.save('data/custom_data/train_data/gray/%d_%d.jpeg'%(num, label))
    else:
        data.save('data/custom_data/train_data/red/%d_%d.jpeg'%(num, label))

학습시작

import torch
import torch.nn as nn
import torch.nn.functional as F

import torch.optim as optim
from torch.utils.data import DataLoader

import torchvision
import torchvision.transforms as transforms

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

torch.manual_seed(777)
if device == 'cuda':
    torch.cuda.manual_seed_all(777)

trans = transforms.Compose([
    transforms.ToTensor()
])

train_data = torchvision.datasets.ImageFolder(root = 'data/custom_data/train_data',
                                             transform = trans)

data_loader = DataLoader(dataset = train_data,
                         batch_size = 8, 
                         shuffle = True,
                         num_workers = 2
                        )

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 6, kernel_size = 5, stride = 1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(6, 16, kernel_size = 5, stride = 1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.layer3 = nn.Sequential(
            nn.Linear(16*13*29, 120),
            nn.ReLU(),
            nn.Linear(120, 2)
        )
        
    def forward(self, x):
        out = self.layer1(x)
#         print(out.shape)  ## 확인하고 지우기
        out = self.layer2(out)
#         print(out.shape)
        out = out.view(out.shape[0], -1)
#         print(out.shape)
        out = self.layer3(out)
        
        return out

# test code
## 클래스가 잘 만들어졌는지 데이터 shape 크기를 확인하는 테스트 코드

net = CNN().to(device) 
test_input = (torch.Tensor(3, 3, 64, 128)).to(device)
test_out = net(test_input)

>>>
torch.Size([3, 6, 30, 62])
torch.Size([3, 16, 13, 29])
torch.Size([3, 6032])

optimizer = optim.Adam(net.parameters(), lr = 0.00001)
loss_func = nn.CrossEntropyLoss().to(device)

total_batch = len(data_loader)

epochs = 7
for epoch in range(epochs):
    avg_cost = 0.0
    for num, data in enumerate(data_loader):
        imgs, labels = data
        imgs = imgs.to(device)
        labels = labels.to(device)
        
        optimizer.zero_grad()
        out = net(imgs)
        loss = loss_func(out, labels)
        loss.backward()
        optimizer.step()
        
        avg_cost += loss / total_batch
        
    print('[Epoch:{}] cost = {}'.format(epoch+1, avg_cost))
print('Learning Finished!')

>>>

[Epoch:1] cost = 0.5202767848968506
[Epoch:2] cost = 0.4274187386035919
[Epoch:3] cost = 0.33061501383781433
[Epoch:4] cost = 0.2468007504940033
[Epoch:5] cost = 0.18349528312683105
[Epoch:6] cost = 0.13799047470092773
[Epoch:7] cost = 0.10548725724220276
Learning Finished!

모델 저장

## model save
torch.save(net.state_dict(), "./model/model.pth")  ## state_dict 저장

# 새로운 모델 생성
new_net = CNN().to(device)

## 새로운 모델에 저장한 모델 statedict 불러와 담기
new_net.load_state_dict(torch.load('./model/model.pth'))

## 기존 모델과 불러온 모델이 동일한가?

print(net.layer1[0])
print(new_net.layer1[0])

print(net.layer1[0].weight[0][0][0])
print(new_net.layer1[0].weight[0][0][0])

net.layer1[0].weight[0] == new_net.layer1[0].weight[0]

>>>
Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
tensor([-0.0186,  0.0603,  0.0351,  0.0412, -0.0270], grad_fn=<SelectBackward>)
tensor([-0.0186,  0.0603,  0.0351,  0.0412, -0.0270], grad_fn=<SelectBackward>)

 

test 시작 

## test data
trans=torchvision.transforms.Compose([
    transforms.Resize((64,128)),
    transforms.ToTensor()
])
test_data = torchvision.datasets.ImageFolder(root='./data/custom_data/test_data', transform=trans)

test_set = DataLoader(dataset = test_data, batch_size = len(test_data))

with torch.no_grad():
    for num, data in enumerate(test_set):
        imgs, label = data
        imgs = imgs.to(device)
        label = label.to(device)
        
        prediction = net(imgs)
        
        correct_prediction = torch.argmax(prediction, 1) == label
        
        accuracy = correct_prediction.float().mean()
        print('Accuracy:', accuracy.item())
        
>>>

Accuracy: 1.0