【PyTorch】AlexNetを使ってCIFAR-10を画像分類してみる

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
from matplotlib import pyplot as plt
%matplotlib inline

train_dataset = torchvision.datasets.CIFAR10(root='./data/',train=True,transform=transforms.ToTensor(),download=True)
test_dataset = torchvision.datasets.CIFAR10(root='./data/',train=False,transform=transforms.ToTensor(),download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,batch_size=64,shuffle=True,num_workers=2)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,batch_size=64,shuffle=False,num_workers=2)

class AlexNet(nn.Module):
  def __init__(self,num_classes=10):
    super(AlexNet,self).__init__()
    self.features = nn.Sequential(
        nn.Conv2d(3,64,kernel_size=11,stride=4,padding=5),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2,stride=2),
        nn.Conv2d(64,192,kernel_size=5,padding=2),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2,stride=2),
        nn.Conv2d(192,384,kernel_size=3,padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(384,256,kernel_size=3,padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(256,256,kernel_size=3,padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(kernel_size=2,stride=2),
    )
    self.classifier = nn.Linear(256,num_classes)

  def forward(self,x):
    x = self.features(x)
    # xを一次元化する
    x = x.view(x.size(0),-1)
    x = self.classifier(x)
    return x
    
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = AlexNet().to(device)

#交差エントロピー誤差を使用
criterion = nn.CrossEntropyLoss()
#モーメンタム法を使用
optimizer = optim.SGD(net.parameters(),lr=0.01,momentum=0.9,weight_decay=5e-4)

# 繰り返し回数
num_epochs = 20
train_loss_list = []
train_acc_list = []
val_loss_list = []
val_acc_list = []

for epoch in range(num_epochs):
  # 値の初期化
  train_loss = 0
  train_acc = 0
  val_loss = 0
  val_acc = 0
  # ここから訓練
  net.train()
  for i,(images,labels) in enumerate(train_loader):
    images,labels = images.to(device),labels.to(device)
    # 勾配を0にする
    optimizer.zero_grad()
    # 順伝播の計算
    outputs = net(images)
    # 誤差の計算
    loss = criterion(outputs,labels)
    # 訓練誤差
    train_loss += loss.item()
    # 予測結果と正解ラベルが同じ数(同じ場合に1となる)
    train_acc += (outputs.max(1)[1]==labels).sum().item()
    # 誤差の逆伝播
    loss.backward()
    # 重みの更新
    optimizer.step()
  # 訓練誤差の平均
  avg_train_loss = train_loss / len(train_loader.dataset)
  # 正解率(訓練)の平均
  avg_train_acc = train_acc / len(train_loader.dataset)

  # ここから評価(dropoutやバッチ正規化を定義した場合に必要)
  net.eval()
  # 学習モデルを固定
  with torch.no_grad():
    for i,(images,labels) in enumerate(test_loader):
      images,labels = images.to(device),labels.to(device)
      # 順伝播の計算
      outputs = net(images)
      # 誤差の計算
      loss = criterion(outputs,labels)
      # 評価誤差
      val_loss += loss.item()
      # 正解率(評価)
      val_acc += (outputs.max(1)[1]==labels).sum().item()
    # 評価誤差の平均
    avg_val_loss = val_loss / len(test_loader.dataset)
    # 正解率(評価)の平均
    avg_val_acc = val_acc / len(test_loader.dataset)

    # プロット用リストへの格納
    train_loss_list.append(avg_train_loss)
    train_acc_list.append(avg_train_acc)
    val_loss_list.append(avg_val_loss)
    val_acc_list.append(avg_val_acc)

plt.figure()
plt.plot(range(num_epochs),train_loss_list,label='train_loss')
plt.plot(range(num_epochs),val_loss_list,label='val_loss')
plt.legend()
plt.xlabel('epoch')
plt.ylabel('loss')
plt.title('Training and Validation loss')
plt.grid()

plt.figure()
plt.plot(range(num_epochs),train_acc_list,label='train_acc')
plt.plot(range(num_epochs),val_acc_list,label='val_acc')
plt.legend()
plt.xlabel('epoch')
plt.ylabel('acc')
plt.title('Training and Validation accuracy')
plt.grid()