[PyTorch 논문 구현] AlexNet : ImageNet Classification with Deep Convolutional Neural Networks

4 분 소요

PyTorch를 사용하여 AlexNet을 구현해보았다.
논문 리뷰는 여기👀🪄

| 모델 구조 파악하기

구현하기에 앞서 논문에서 언급한 내용을 바탕으로 레이어 별 구조를 재구성해보았다.

Layer	Kernel (개수)	Activation fn.	Normalization	Feature map (output)
Input	-	-	-	227x227x3
Conv1	11x11x3 (96), stride 4	ReLU	LRN	55x55x96
MaxPool	3x3, stride 2	-	-	27x27x96
Conv2	5x5x48 (256), stride 1, padding 2	ReLU	LRN	27x27x256
MaxPool	3x3, stride 2	-	-	27x27x96
Conv3	3x3x256 (384), stride 1, padding 1	ReLU	-	13x13x384
Conv4	3x3x192 (384), stride 1, padding 1	ReLU	-	13x13x384
Conv5	3x3x192 (256), stride 1, padding 1	ReLU	-	13x13x256
MaxPool	3x3, stride 2	-	-	6x6x256
FC1	-	ReLU	-	4096
FC2	-	ReLU	-	4096
FC3	-	ReLU	-	1000 (클래스 개수)

Convolution layer의 padding과 stride가 주어지지 않은 경우, 아래의 식을 활용하여 계산했다. 여기서 input과 output은 각각의 feature map 크기를 의미한다. 만약 계산 결과가 정수로 딱 떨어지지 않는다면 버림으로 계산하면 된다.

\[\frac{input + 2 * padding - kernel}{stride}+1=output\]

예를 들어, Conv3의 stride와 padding을 구해보자.
input = 13, output = 13, kernel = 3 이라고 주어졌으므로 이 값들을 식에 대입하고 정리하면 아래와 같다.

\[\frac{13 + 2 * padding - 3}{stride}+1=13\] \[10+2*padding=12*stride\]

깔끔하게 값이 구해지진 않았지만, 직관적으로 padding=1, stride=1이면 식이 성립함을 알 수 있다.

| 데이터 불러오기

논문에선 ImageNet 데이터를 사용했지만, 이번엔 간단하게 CIFAR 10 데이터를 사용하기로 했다. CIFAR-10 데이터셋은 60000장으로 이루어진 32x32 크기의 RGB 이미지 데이터이며 10개의 클래스로 이루어져 있다.

전처리로는 이미지의 크기를 227x227로 맞추는 resize와 데이터 셋의 평균과 표준편차를 활용한 normalization을 적용했다.

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((227,227)),
    transforms.Normalize(mean = [0.4914, 0.4822, 0.4465],
                         std = [0.2023, 0.1994, 0.2010])
])

갑자기 등장한 데이터 셋의 평균(mean)과 표준편차(std) 값은 아래의 코드를 통해 구했다.
코드는 꾸준희님의 티스토리를 참고했다.

#import torch
#import torchvision
#import torchvision.datasets as datasets
#from tqdm import tqdm

dataset = datasets.CIFAR10("./data", download=True, train=True, transform=transforms.ToTensor())
full_loader = torch.utils.data.DataLoader(dataset, shuffle=False)

mean = torch.zeros(3)
std = torch.zeros(3)

for image, _label in tqdm(full_loader) :
  for i in range(3) :
    mean[i] += image[:,i,:,:].mean()
    std[i] += image[:,i,:,:].std()

# div_() = inplace version of div()
mean.div_(len(dataset))
std.div_(len(dataset))

print('\n')
print(f'mean : {mean}')
print(f'std : {std}')

Torch Tensor vs. Numpy ndarray
(H = height, W = weight, C = channel, B = batch size)
torch.tensor: C x H x W (3차원) or B x C x H x W (4차원)
numpy.ndarray : H x W x C

코드를 보다보니까 텐서의 채널이 의미하는 바가 헷갈려서 정리해봤다.
현재 이미지는 4차원 텐서이고, 우리는 채널 별로 mean, std 값을 구해야하므로 텐서의 두 번째 인덱스에 대해 연산을 수행했다.

이제 train, test set을 불러오고 data loader에 넣어주면 된다. 짠!

#import torchvision
#import torchvision.datasets as datasets
#import torchvision.transforms as transforms

train = datasets.CIFAR10(root='./data', train=True, transform=transform, download=False)
test = datasets.CIFAR10(root='./data', train=False, transform=transform, download=False)

print('train: \n', train, '\n')
print('test: \n', test, '\n')

BATCH_SIZE = 128

train_loader = torch.utils.data.DataLoader(dataset=train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test, batch_size=BATCH_SIZE, shuffle=True)

| 모델 구현하기

위에서 작성한 표 그대로! 레이어를 쌓아주자.

class AlexNet(nn.Module) :

  """
    AlexNet Class
  """

  def __init__(self, num_classes) :
    super().__init__()

    self.num_classes = num_classes

    self.conv1 = nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4),
        nn.ReLU(),
        nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
        nn.MaxPool2d(kernel_size=3, stride=2)
    )

    self.conv2 = nn.Sequential(
        nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2),
        nn.ReLU(),
        nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
        nn.MaxPool2d(kernel_size=3, stride=2)
    )

    self.conv3 = nn.Sequential(
        nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1),
        nn.ReLU()
    )

    self.conv4 = nn.Sequential(
        nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
        nn.ReLU()
    )

    self.conv5 = nn.Sequential(
        nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=3, stride=2)
    )

    self.fc1 = nn.Sequential(
        nn.Dropout(),
        nn.Linear(in_features=6*6*256, out_features=4096),
        nn.ReLU()
    )

    self.fc2 = nn.Sequential(
        nn.Dropout(),
        nn.Linear(in_features=4096,out_features=4096),
        nn.ReLU()
    )

    self.fc3 = nn.Linear(in_features=4096, out_features=self.num_classes)

  def forward(self, x) :
    # convolution laters
    x = self.conv1(x)
    x = self.conv2(x)
    x = self.conv3(x)
    x = self.conv4(x)
    x = self.conv5(x)

    # fc layers
    x = x.view(-1, 256*6*6) # 이거 안해주면 오류나요 (flatten)
    x = self.fc1(x)
    x = self.fc2(x)
    x = self.fc3(x)
    return x

처음에 텐서를 1차원으로 펴주는 과정 없이 모델 구현을 끝냈었다. 그런데 모델을 돌리려고 하니까 아래의 오류가 발생하는거 아닌가 (•᷄ࡇ•᷅) 과장해서 한 시간 정도 찾아보다가 겨우 깨닫고 해결했다.
잊지 말자. FC layer 전에 텐서를 1차원으로 만들어줘야 한다.

| 모델 확인하기

이제 torchsummary로 모델을 잘 작성했는지 확인해보면 !

import torchsummary

NUM_CLASSES = 10

model = AlexNet(NUM_CLASSES)
model.to(device)

torchsummary.summary(model, input_size=(3,227,227), device='cuda')

완벽하다.

| 모델 학습하기

데이터도 불러 왔고, 모델도 구현했으니까 이제 학습만 하면 된다.
일단 해결하려는 문제는 분류 문제이므로 손실 함수로는 Cross Entropy Loss를 사용했고, Optimizer는 논문에서 언급한 그대로 설정했다. Epoch는 내 맘대로 10으로 설정했다.

EPOCHS = 10
criterion = nn.CrossEntropyLoss()
optimizer = SGD(model.parameters(), lr=0.0001, momentum=0.9, weight_decay=0.0005)

def model_train(model, train_loader, criterion, optimizer, device) :
  model.train()
  for epoch in range(EPOCHS) :
    train_loss_sum = 0
    train_n_total, train_n_corr = 0, 0

    for iter, (img, label) in enumerate((tqdm(train_loader))) :
      img, label = img.to(device), label.to(device)

      # gradient reset
      optimizer.zero_grad()

      # forward pass
      pred = model.forward(img)

      # calculate loss
      loss = criterion(pred, label)

      # backpropagate, calculate derivative
      loss.backward()

      # optimizer update, gradient update
      optimizer.step()

      train_loss_sum += loss.item()

      _, y_pred = pred.max(dim=1)
      train_n_total += img.size(0)
      train_n_corr += (y_pred == label).sum().item()

    train_loss_avg = train_loss_sum/len(train_loader)
    train_acc = train_n_corr/train_n_total

    print(f'Epoch: {epoch+1}/{EPOCHS}, train loss : {train_loss_avg}, train acc : {train_acc}')

  print('Done')

zero_grad() ?
PyTorch는 backpropagation 단계에서 gradient를 누적하여 더한다. 그러나 파라미터는 현재 미분값을 이용하여 없데이트 된다. 누적된 미분값을 파라미터 업데이트에 사용한다면 분명히 오류가 발생할 것이다. 이를 방지하기 위해 zero_grad()로 매 iter마다 gradient를 초기화 해준다.

이제 학습을 돌려보자.

model_train(model, train_loader, criterion, optimizer, device)

오… 성능이 너무 별로여서 깜짝 놀랐다 ꒰⍤꒱ 에포크 수가 너무 적었던 걸까?
조금 당황스러웠지만 그래도 구현을 해봤다는 점에 의의를 두기로 했다.

| 모델 테스트하기

마지막으로 학습한 모델을 테스트해보자.

def model_test(model, test_loader, criterion, device) :
  model.eval()
  with torch.no_grad() :
    loss_sum = 0
    n_total, n_corr = 0, 0

    for iter, (img, label) in enumerate(tqdm(test_loader)) :
      img, label = img.to(device), label.to(device)

      pred = model.forward(img)

      _, y_pred = pred.max(dim=1)
      n_corr += (label == y_pred).sum().item()
      n_total += img.size(0)

      loss = criterion(pred, label)
      loss_sum += loss

    loss_avg = loss_sum/len(test_loader)
    acc = n_corr/n_total

    print(f'\nLoss : {loss_avg}, Acc : {acc}')

model_test(model, test_loader, criterion, optimizer, device)

역시 정확도가 아쉽다.
이렇게 AlexNet 구현 및 학습 끝!

Twitter Facebook LinkedIn

codehyunn

[PyTorch 논문 구현] AlexNet : ImageNet Classification with Deep Convolutional Neural Networks

| 모델 구조 파악하기

| 데이터 불러오기

| 모델 구현하기

| 모델 확인하기

| 모델 학습하기

| 모델 테스트하기

공유하기

댓글남기기

참고

[논문 리뷰] LoRA: Low Rank Adaptation of Large Language Models

[논문 리뷰] Attention is all you need

[NLP] Beam Search, BLEU Score

[NLP] Seq2Seq with Attention