[CNN] Alex Net - 구현 및 MNIST_Fashion 학습

AlexNet을 모듈로 표현해보기 위해 구현하였으며 이미지 또한 AlexNet에 맞춰 사용이 가능하도록 구현하였다.

AlexNet은 2개의 GPU로 병렬 연산한다는게 포인트이지만, 그래픽카드가 하나인 관계로 병렬연산은 구현하지 못했다. 

---

AlexNet 구현

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms


class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4)
        self.conv2 = nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2)
        self.conv3 = nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(in_features=256 * 6 * 6, out_features=4096)
        self.fc2 = nn.Linear(in_features=4096, out_features=4096)
        self.fc3 = nn.Linear(in_features=4096, out_features=10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=3, stride=2)

        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, kernel_size=3, stride=2)

        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = F.relu(self.conv5(x))
        x = F.max_pool2d(x, kernel_size=3, stride=2)

        x = x.view(x.size(0), -1)

        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=0.5)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, p=0.5)
        x = F.log_softmax(self.fc3(x), dim=1)

        return x

    @staticmethod
    def transform():
        return transforms.Compose([transforms.Resize((227, 227)),
                                   transforms.ToTensor(),
                                   transforms.Normalize(mean=(0.1307,), std=(0.3081,))])


if __name__ == "__main__":
    # if gpu is to be used
    use_cuda = torch.cuda.is_available()
    print("use_cuda : ", use_cuda)

    FloatTensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
    device = torch.device("cuda:0" if use_cuda else "cpu")

    net = AlexNet().to(device)

    X = torch.randn(size=(1, 1, 227, 227)).type(FloatTensor)

    print(net(X))

---

구현한 AlexNet를 import하여 MNIST_Fashion 학습

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets

from CNN.AlexNet import AlexNet

if __name__=="__main__":
    # hyper parameter
    batch_size = 100
    num_epochs = 250
    learning_rate = 0.0001

    # data load
    root = './MNIST_Fashion'
    transform = AlexNet.transform()

    train_set = datasets.FashionMNIST(root=root, train=True, transform=transform, download=True)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True)

    test_set = datasets.FashionMNIST(root=root, train=False, transform=transform, download=True)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=True)

    # if gpu is to be used
    use_cuda = torch.cuda.is_available()
    print("use_cuda : ", use_cuda)

    device = torch.device("cuda:0" if use_cuda else "cpu")

    Net = AlexNet().to(device)

    criterion = nn.CrossEntropyLoss().to(device)
    optimizer = optim.Adam(Net.parameters(), lr=learning_rate)

    # train
    costs = []
    total_batch = len(train_loader)
    for epoch in range(num_epochs):
        total_cost = 0
        for imgs, labels in train_loader:
            imgs, labels = imgs.to(device), labels.to(device)

            predict = Net(imgs)
            loss = criterion(predict, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_cost += loss

        avg_cost = total_cost / total_batch
        print("Epoch:", "%03d" % (epoch + 1), "Cost =", "{:.9f}".format(avg_cost))
        costs.append(avg_cost)