[Pytorch] CNN으로 사진 분류하기

1. 데이터 전처리하기
   • 데이터 증강
   • 이미지 정규화
2. CNN으로 이미지 분류하기
   • 모델 정의하기
   • 모델 학습하기
   • 모델 성능 평가하기
3. 전이 학습 모델 VGG로 분류하기
   • 사전 학습된 모델 불러오기
   • 학습 루프 정의하기
   • 학습 및 성능 평가하기

---

핵심 용어 정리

1. convolution : 작은 필터를 이용해 이미지로부터 특징 추출
2. CNN : convolution layer를 반복적으로 쌓아 만든 인공 신경망
3. feature map : convolution layer의 결과를 말함
4. 데이터 증강 : 이미지의 회전 및 잘라내기 등 하나로 여러가지의 데이터를 만들어 데이터를 늘리는 기법
5. 이미지 정규화 : 이미지의 픽셀 간 편향 제거에 사용. 각 채널의 분포가 동일해짐
6. padding : 이미지 외곽을 0으로 채워서 convolution 전후 크기를 동일하게 만듬
7. cropping : 이미지의 일부를 잘라냄
8. max pooling :  이미지의 크기를 줄이는 데 사용하는 기법으로 커널에서 가장 큰 값을 이용
9. Transfer learning : 사전 학습된 모델의 파라미터를 수정해서 자신의 데이터셋에 최적화시키는 방법. 학습시간 단축 가능

 

 

데이터를 먼저 받아오자 CIFAR10의 데이터를 통해 진행할 것이다.

 

import matplotlib.pyplot as plt

from torchvision.datasets.cifar import CIFAR10
from torchvision.transforms import ToTensor

# CIFAR-10 데이터셋 불러오기
training_data = CIFAR10(
    root="../data/",
    train=True,
    download=True,
    transform=ToTensor())

test_data = CIFAR10(
    root="../data/",
    train=False,
    download=True,
    transform=ToTensor())

for i in range(9):
    plt.subplot(3, 3, i+1)
    plt.imshow(training_data.data[i])
plt.show()

상위에 data폴더를 만들었고 그곳에 데이터를 받았다.

트레이닝 5만개, 테스트 1만개지만 데이터를 좀더 사용하기 위해서 증강을 시켜봅니다. 

import torchvision.transforms as T

from torchvision.transforms import Compose
from torchvision.transforms import RandomHorizontalFlip, RandomCrop

transforms = Compose([  # 데이터전처리함수 순서대로 넣어주면 적용됨
    T.ToPILImage(),
    RandomCrop((32, 32), padding=4), # 랜덤으로 이미지 일부 제거 후 패딩
    RandomHorizontalFlip(p=0.5),    # y축을 기준으로 대칭
])

training_data = CIFAR10(
    root="../data/",
    train=True,
    download=True,
    transform=ToTensor())

test_data = CIFAR10(
    root="../data/",
    train=False,
    download=True,
    transform=ToTensor())

for i in range(9):
    plt.subplot(3, 3, i+1)
    plt.imshow(transforms(training_data.data[i]))
plt.show()

 

데이터가 일부 잘리거나 좌우대칭으로 바뀐 것을 볼 수 있습니다. 이로써 다양성을 높입니다

 

또한, rgb컬러의 값의 분포를 정규화하여 학습이 용이하게 일어나게 합니다.

 

이미지들의 rgb 평균과 표준편차를 구하는 코드입니다.

 

import torch 

training_data = CIFAR10(
    root="../data/",
    train=True,
    download=True,
    transform=ToTensor())

test_data = CIFAR10(
    root="../data/",
    train=False,
    download=True,
    transform=ToTensor())

imgs = [item[0] for item in training_data]

# 이미지를 하나로 합침
imgs = torch.stack(imgs, dim=0).numpy()

# rgb 각 평균
mean_r = imgs[:,0,:,:].mean()
mean_g = imgs[:,1,:,:].mean()
mean_b = imgs[:,2,:,:].mean()
print(mean_r, mean_g, mean_b)

# rgb 각 표준편차
std_r = imgs[:,0,:,:].std()
std_g = imgs[:,1,:,:].std()
std_b = imgs[:,2,:,:].std()
print(std_r, std_g, std_b)

 

이를 이용하여 정규화를 진행하면 다음과 같이 코드를 작성할 수 있습니다.

import torchvision.transforms as T

from torchvision.transforms import Compose
from torchvision.transforms import RandomHorizontalFlip, RandomCrop, Normalize

transforms = Compose([  # 데이터전처리함수 순서대로 넣어주면 적용됨
    T.ToPILImage(),
    RandomCrop((32, 32), padding=4), # 랜덤으로 이미지 일부 제거 후 패딩
    RandomHorizontalFlip(p=0.5),    # y축을 기준으로 대칭
    T.ToTensor(),
    
    Normalize(mean=(mean_r, mean_g, mean_b), std=(std_r, std_g, std_b)),
    T.ToPILImage()
])

training_data = CIFAR10(
    root="../data/",
    train=True,
    download=True,
    transform=ToTensor())

test_data = CIFAR10(
    root="../data/",
    train=False,
    download=True,
    transform=ToTensor())

for i in range(9):
    plt.subplot(3, 3, i+1)
    plt.imshow(transforms(training_data.data[i]))
plt.show()

 

 

 

 

다음은 CNN학습을 위한 모델을 구성하는 코드를 작성한다

import torch
import torch.nn as nn

class BasicBlock(nn.Module):
    # 기본 블럭을 구성하는 층 정의
    def __init__(self, in_channels, out_channels, hidden_dim):
        # nn.Module 클래스의 요소 상속
        super(BasicBlock, self).__init__()
        # convolution layer 정의
        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(hidden_dim, out_channels, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        
        # stride는 커널의 이동 거리
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
    
    def forward(self, x): # 기본 블럭의 순전파 정의
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool(x)
        return x

convolution layer를 정의한다. 각 layer를 통과하고 relu를 통과시켜준다.

 

class CNN(nn.Module):
    def __init__(self, num_classes): # 클래스의 갯수
        super(CNN, self).__init__()
        
        # 합성곱 기본 블럭 정의
        self.block1 = BasicBlock(in_channels=3, out_channels=32, hidden_dim=16)
        self.block2 = BasicBlock(in_channels=32, out_channels=128, hidden_dim=64)
        self.block3 = BasicBlock(in_channels=128, out_channels=256, hidden_dim=128)
        
        # 분류기 정의
        self.fc1 = nn.Linear(in_features=4096, out_features=2048)
        self.fc2 = nn.Linear(in_features=2048, out_features=256)
        self.fc3 = nn.Linear(in_features=256, out_features=num_classes)
        
        # 분류기의 활성화 함수
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x) # 출력모양: (-1, 256, 4, 4)
        x = torch.flatten(x, start_dim=1) # 2차원의 feature map을 1차원으로
        
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        
        return x

Linear를 통한 선형분류작업을 진행한다.

 

 

 

자 이제 모델 학습을 진행해보자

from torch.utils.data.dataloader import DataLoader

from torch.optim.adam import Adam

transforms = Compose([
    RandomCrop((32, 32), padding=4), # 랜덤 Cropping
    RandomHorizontalFlip(p=0.5), # y축으로 좌우대칭
    ToTensor(), # 텐서로 변환
    Normalize(mean=(mean_r, mean_g, mean_b), std=(std_r, std_g, std_b)) # meand rgb: 0.49139968 0.48215827 0.44653124, std rgb: 0.24703233 0.24348505 0.26158768
])

# 데이터 불러오기
training_data = CIFAR10(root="../data/", train=True, download=True, transform=transforms)
test_data = CIFAR10(root="../data/", train=False, download=True, transform=transforms)

# 데이터로더 정의
train_loader = DataLoader(training_data, batch_size=32, shuffle=True)
test_loader = DataLoader(test_data, batch_size=32, shuffle=False)

# 학습을 진행할 프로세서 설정
device = "cuda" if torch.cuda.is_available() else "cpu"

## 맥일 경우 아래코드 사용
# device = torch.device('mps:0' if torch.backends.mps.is_available() else 'cpu')

print(f"device 정보: {device}")
# CNN 모델 정의
model = CNN(num_classes=10)

# 모델을 device로 보냄
model.to(device)

맥일 경우의 device코드도 적어두었으니 참고!

 

이렇게 만든 모델이 정리가 되었다.

 

# 학습률 정의
lr = 1e-3

# 최적화 기법 정의
optim = Adam(model.parameters(), lr=lr)

# 학습 루프 정의
for epoch in range(100):
    for data, label in train_loader: # 데이터 호출
        optim.zero_grad() # 기울기를 0으로 초기화
        
        preds = model(data.to(device)) # 모델의 예측
        
        # 오차역전파와 최적화 (손실을 계산하고 backpropagation을 통해 최적화함)
        loss = nn.CrossEntropyLoss()(preds, label.to(device))
        loss.backward()
        optim.step()
    
    if epoch==0 or epoch%10==9: # 10번마다 손실 출력
        print(f"epoch{epoch+1} loss:{loss.item()}")
        
# 모델 저장
torch.save(model.state_dict(), "CIFAR.pth")

100번의 에포크를 주고 10번마다 loss값을 출력한다.

 

이후 모델을 저장한다.

 

나는 gtx1060 6gb모델이라 그런지 꽤나 오래걸렸다.

55분정도 소요!!

 

 

 

model.load_state_dict(torch.load("CIFAR.pth", map_location=device))

num_corr = 0

with torch.no_grad():
    for data, label in test_loader:
        output = model(data.to(device))
        preds = output.data.max(1)[1]
        corr = preds.eq(label.to(device).data).sum().item()
        num_corr += corr

    print(f"Accuracy:{num_corr/len(test_data)}")

Accuracy:0.8305

모델성능은 83%정도로 쏘쏘했다. layer를 적당하게 넣었기 때문에 아주 훌륭한 정확도가 뜨진 않은 것 같다.

 

 

 

이렇게  레이어를 쌓아서 만들어도되지만 기존에 있던 모델을 사용할 수도 있다.

전이학습형태로!

ImageNet으로 사전학습한 vgg16 모델을 사용해보자

import torch
import torch.nn as nn
from torchvision.models.vgg import vgg16
 
device = "cuda" if torch.cuda.is_available() else "cpu"
 
model = vgg16(pretrained=True) # 모델 객체 생성
fc = nn.Sequential( # 분류층 정의
                    nn.Linear(512 * 7 * 7, 4096),
                    nn.ReLU(),
                    nn.Dropout(),
                    nn.Linear(4096,4096),
                    nn.ReLU(),
                    nn.Dropout(),
                    nn.Linear(4096, 10),
                    )

model.classifier = fc 
model.to(device)

 

이후 데이터 처리나 이런것들은 동일하게 가져가자 피처값을 잘보고 가면된다

 

import tqdm

from torchvision.datasets.cifar import CIFAR10
from torchvision.transforms import Compose, ToTensor, Resize
from torchvision.transforms import RandomHorizontalFlip, RandomCrop, Normalize
from torch.utils.data.dataloader import DataLoader

from torch.optim.adam import Adam

transforms = Compose([
    Resize(224),
    RandomCrop((224, 224), padding=4),
    RandomHorizontalFlip(p=0.5),
    ToTensor(),
    Normalize(mean=(mean_r, mean_g, mean_b), std=(std_r, std_g, std_b))
])

# 데이터 불러오기
training_data = CIFAR10(root="../data/", train=True, download=False, transform=transforms)
test_data = CIFAR10(root="../data/", train=False, download=False, transform=transforms)

# 데이터로더 정의
train_loader = DataLoader(training_data, batch_size=32, shuffle=True)
test_loader = DataLoader(test_data, batch_size=32, shuffle=False)

 

lr = 1e-4
optim = Adam(model.parameters(), lr=lr)

for epoch in range(30):
    iterator = tqdm.tqdm(train_loader) # 학습로그 출력
    for data, label in iterator:
        optim.zero_grad()
        preds = model(data.to(device)) # 모델의 예측값 출력
        
        loss = nn.CrossEntropyLoss()(preds, label.to(device))
        loss.backward()
        optim.step()
        
        # tqdm이 출력할 문자열
        iterator.set_description(f"epoch:{epoch+1} loss:{loss.item()}")
        
torch.save(model.state_dict(), "CIFAR_pretrained.pth") # 모델저장

참고로 kaggle p100환경에서 돌려본 결과 epoch한번에 1시간가까이 돌아갔다

 

model.load_state_dict(torch.load("CIFAR_pretrained.pth", map_location=device))

num_corr = 0

with torch.no_grad():
    for data, label in test_loader:
        
        output = model(data.to(device))
        preds = output.data.max(1)[1]
        corr = preds.eq(label.to(device).data).sum().item()
        num_corr += corr
        
    print(f"Accuracy:{num_coprr/len(test_data)}")