Pytorch를 사용하여 CIFAR10 이미지 분류기 만들기

# Package import
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torchvision import transforms

import matplotlib.pyplot as plt
import numpy as np
import random

# Set the random seeds for reproducibility
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

transform = transforms.Compose(
    # numpy array -> Torch Tensor로 변환
    [transforms.ToTensor(),
     # RGB 채널에 0.5를 뺀 후 표준편차를 0.5로 나누어 이미지 정규화
     transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))]
)

# batch size 정의
batch_size = 128

# 학습 데이터 다운
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
# 테스트 데이터 다운
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 학습, 검증 데이터로더 정의
trainset_size = int(0.8 * len(trainset))
valset_size = len(trainset) - trainset_size
trainset, valset = torch.utils.data.random_split(trainset, [trainset_size, valset_size])

trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True)
valloader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False)

# 테스트 데이터로더 정의
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False)

train_samples = len(trainset)
val_samples = len(valset)
test_samples = len(testset)
image_shape = testset.data.shape[1:]

print(f'Number of Train samples: {train_samples}')
print(f'Number of Val samples: {val_samples}')
print(f'Number of Test samples: {test_samples}')
print(f'Image shape: {image_shape}')

# 이미지 확인을 위한 함수 정의
def imshow(img):
    # 정규화 해제
    img = img / 2 + 0.5 
    npimg = img.numpy()
    # npimg 축 순서 변경
    plt.figure(figsize=(12,18))
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 이미지 가져오기
dataiter = iter(trainloader)
images, labels = next(dataiter)

# 이미지 보여주기
# images 변수에 저장된 이미지들을 그리드 형태로 만들어 시각화하는 코드
imshow(torchvision.utils.make_grid(images))

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# hyper Parameter
learning_rate = 0.001
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Loss Function
mymodel = MyModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(mymodel.parameters(), lr=learning_rate)

num_epochs = 10
for epoch in range(num_epochs): 
a    running_loss = 0.0

    # Simple CNN 학습
    mymodel.train()
    for i, data in enumerate(trainloader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = mymodel(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(trainloader)}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0

    # Simple CNN 검증
    mymodel.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for inputs, labels in valloader:
            inputs, labels = inputs.to(device), labels.to(device)

            outputs = mymodel(inputs)
            _, predicted = torch.max(outputs.data, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print(f'Epoch [{epoch+1}/{num_epochs}], Accuracy of the network on the Val images: {100 * correct / total}%')

print('Finished Training')

# 정확도 확인
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = mymodel(images)
        
        # 가장 높은 값를 갖는 분류(class)를 정답으로
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')

# 각 분류(class)에 대한 예측값 계산을 위해 class명 정의
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = mymodel(images)
        _, predictions = torch.max(outputs, 1)
        # 각 분류별로 올바른 예측 수 저장
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# 각 분류별 정확도(accuracy)를 출력
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')

# Load the pre-trained VGG16 model
vgg16 = torchvision.models.vgg16(pretrained=True)

# 마지막 Layer를 CIFAR-10 데이터에 맞게 재정의
num_classes = 10
vgg16.classifier[6] = torch.nn.Linear(vgg16.classifier[6].in_features, num_classes)

# HypterParameter 정의
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(vgg16.parameters(), lr=learning_rate)

# VGG16 Fine Tunning
vgg16.to(device)
for epoch in range(num_epochs):
    running_loss = 0.0
    vgg16.train()
    for i, (inputs, labels) in enumerate(trainloader):
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = vgg16(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(trainloader)}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0
            
    # VGG16 검증
    vgg16.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for inputs, labels in valloader:
            inputs, labels = inputs.to(device), labels.to(device)

            outputs = vgg16(inputs)
            _, predicted = torch.max(outputs.data, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print(f'Accuracy of the network on the Val images: {100 * correct / total}%')

print('Finished Training')

# Evaluate the model
vgg16.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for inputs, labels in testloader:
        inputs, labels = inputs.to(device), labels.to(device)

        outputs = vgg16(inputs)
        _, predicted = torch.max(outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network on the test images: {100 * correct / total}%')

# 각 분류(class)에 대한 예측값 계산을 위해 class명 정의
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = vgg16(images)
        _, predictions = torch.max(outputs, 1)
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# 각 분류별 정확도(accuracy)를 출력합니다
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')