pytorch-5 CNN

1. CNN

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "True"

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


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)

        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2)

        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)

        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2)

        self.fc = nn.Linear(32 * 7 * 7, 10)

        self.conv1_output = None
        self.conv2_output = None

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        self.conv1_output = x  

        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        self.conv2_output = x  

        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


if __name__ == "__main__":
    train_dataset = MNIST(
        "./train", train=True, transform=transforms.ToTensor(), download=True
    )
    test_dataset = MNIST("./train", train=False, transform=transforms.ToTensor())

    batch_size = 64
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)

    model = CNN().to(device)

    criterion = nn.CrossEntropyLoss()

    optimizer = optim.Adam(model.parameters(), lr=0.001)

    num_epochs = 10

    fig, axs = plt.subplots(2, 3, figsize=(10, 8))
    fig.tight_layout(pad=4.0)
    axs = axs.flatten()

    epoch_losses = []

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0

        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()  

            outputs = model(images)

            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

            running_loss += loss.item() * images.size(0)

        epoch_loss = running_loss / len(train_dataset)
        epoch_losses.append(epoch_loss)

        print(f"Epoch {epoch+1} / {num_epochs}, Loss : {epoch_loss:.4f}")

        if epoch == 0:
            weights = model.conv1.weight.detach().cpu().numpy()
            axs[0].imshow(weights[0, 0], cmap="coolwarm")
            axs[0].set_title("Conv1 Weights")
            divider = make_axes_locatable(axs[0])
            cax = divider.append_axes("right", size="5%", pad=0.05)
            plt.colorbar(axs[0].imshow(weights[0, 0], cmap="coolwarm"), cax=cax)

            weights = model.conv2.weight.detach().cpu().numpy()
            axs[1].imshow(weights[0, 0], cmap="coolwarm")
            axs[1].set_title("Conv2 Weights")
            divider = make_axes_locatable(axs[1])
            cax = divider.append_axes("right", size="5%", pad=0.05)
            plt.colorbar(axs[1].imshow(weights[0, 0], cmap="coolwarm"), cax=cax)


            if model.conv1_output is not None:
                conv1_output = model.conv1_output.detach().cpu().numpy()
                axs[2].imshow(weights[0, 0], cmap="coolwarm")
                axs[2].set_title("Conv1 Outputs")
                divider = make_axes_locatable(axs[2])
                cax = divider.append_axes("right", size="5%", pad=0.05)
                plt.colorbar(
                    axs[2].imshow(conv1_output[0, 0], cmap="coolwarm"), cax=cax
                )

            if model.conv2_output is not None:
                conv2_output = model.conv2_output.detach().cpu().numpy()
                axs[3].imshow(weights[0, 0], cmap="coolwarm")
                axs[3].set_title("Conv2 Outputs")
                divider = make_axes_locatable(axs[3])
                cax = divider.append_axes("right", size="5%", pad=0.05)
                plt.colorbar(
                    axs[3].imshow(conv2_output[0, 0], cmap="coolwarm"), cax=cax
                )

    axs[4].plot(range(epoch + 1), epoch_losses)
    axs[4].set_title("training_loss")
    axs[4].set_xlabel("Epoch")
    axs[4].set_ylabel("Loss")

    plt.show()

    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)

            temp, predicted = torch.max(outputs.data, 1)

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

    accuracy = (correct / total) * 100
    print(f"Test Accuracy: {accuracy:.2f} %")
    
Epoch 1 / 10, Loss : 0.2587
Epoch 2 / 10, Loss : 0.0704
Epoch 3 / 10, Loss : 0.0524
Epoch 4 / 10, Loss : 0.0423
Epoch 5 / 10, Loss : 0.0361
Epoch 6 / 10, Loss : 0.0303
Epoch 7 / 10, Loss : 0.0266
Epoch 8 / 10, Loss : 0.0236
Epoch 9 / 10, Loss : 0.0208
Epoch 10 / 10, Loss : 0.0170
Test Accuracy: 98.92%

간단한 CNN모델로 mnist데이터셋을 학습하고 평가하는 코드이다.

 

init 함수에서 각 신경망의 변숫값을 지정하여 최종적으로 10개의 클래스를 분류할 수 있도록 구성하고 forward 함수에서 conv, pooling, full connected 망까지 거쳐서 결괏값을 반환하는 클래스를 만든다.

 

이후에 mnist 데이터셋을 불러와 학습데이터와 테스트 데이터로 쪼개주고 손실함수, 최적화알고리즘을 정해준다. 시각화를 위한 변수들을 작성하고 2차원데이터를 flatten 해서 시각화할 수 있게 한다.

 

이후에 반복문을 돌면서 학습 데이터셋에서 데이터와 라벨값을 꺼내서 순전파, 역전파, 재조정의 과정을 진행하며 각 배치의 손실들을 누적하여 epoch_losses에 리스트로 만든다.

epoch가 0일 때, 첫 번째 반복문에서의 conv1,2의 가중치값, conv1,2의 출력값을 시각화하며 에포크들의 손실값의 변화를 그래프로 확인한다.

 

이후에 모델을 평가하기 위해서 model.eval()으로 평가 모드로 전환하고 테스트 데이터의 개수와 정확하게 맞춘 개수를 구하기 위한 변수를 할당한다.

 

with torch.no_grad()는 그래디언트 계산이 필요 없으므로 비활성화하여 메모리를 향상한다. 테스트 데이터셋에서 데이터와 라벨값을 가져와 모델을 돌리고 그중에서 가장 큰 값을 가져와서 예측된 라벨로 사용한다.

 

테스트 데이터의 총개수를 total에 할당하고 예측된 라벨과 테스터 데이터셋의 라벨값을 비교하여 그 개수를 correct에 저장한다. 이후에 정확도를 계산해서 출력한다. 98프로로 상당히 높은 적중률을 가지고 손글씨 데이터를 분류한다는 것을 알 수 있다.