PyTorchで学ぶディープラーニング：はじめての画像分類 -実装編-

# 必要なライブラリをインポート
import collections                       # ラベルの頻出頻度を調べる

import torch                             # PyTorch全般
import torch.nn as nn                    # ネットワークの構築
import torch.optim as optim              # 学習の最適化を
from torch.utils.data import DataLoader  # データを効率よく呼び出す

import torchvision.transforms as transforms # 画像の正規化などの前処理
import torchvision.utils as utils           # 画像表示用のユーティリティ
from torchvision.datasets import LFWPeople  # LFWデータセット

import matplotlib.pylab as plt   # データのプロット
import numpy as np　              # 画像配列を数値配列として扱う

print(torch.__version__)
torch.manual_seed(0)

transform = transforms.Compose([
    transforms.Resize((64, 64)),                # 画像のリサイズ
    transforms.ToTensor(),                      # PyTorchのテンソル形式へ変換
    transforms.Normalize(mean=(0.5, 0.5, 0.5),
                         std=(0.5, 0.5, 0.5))   # RGB各チャンネルの平均・標準偏差を正規化
])

lfw_dataset = LFWPeople(root='./data', download=True, transform=transform)

# データセットからラベルの抜き出し
label_list = lfw_dataset.targets
print(label_list[:10])   # データ確認のため最初の10個表示

# ラベルの出現頻度別に(ラベル, 出現数)で抜き出し
label_counter = collections.Counter(label_list)
freq_list = label_counter.most_common()
print(freq_list[:10])    # データ確認のため最初の10個表示

# 上位num_classes分のラベルを抜き出し
num_classes = 10
top_freq_label = [label for label, _ in freq_list[:num_classes]] # リストの最初にラベルが入っている
print(top_freq_label)

# ラベルを0から採番し直す
# トップ10に入っている画像と採番し直したラベルを新たなデータセットに格納する
label_mapping = {label: i for i, label in enumerate(top_freq_label)}
top_freq_lfw_dataset = [(img, label_mapping[label]) for img, label in lfw_dataset if label in top_freq_label]

# ハイパーパラメータの定義
batch_size = 24

# データセットを学習データセットと検証データセットに分類する
train_size = int(0.8 * len(top_freq_lfw_dataset))
valid_size = len(top_freq_lfw_dataset) - train_size
train_dataset, valid_dataset = torch.utils.data.random_split(top_freq_lfw_dataset, [train_size, valid_size])

# データローダの作成
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

def imshow(img):
    # matplotlibに合うように正規化の修正
    img = img / 2 + 0.5
    # matplotlibで出力できるようnumpyへ変換
    np_img = img.numpy()
    # [C, H, W] -> [H, W, C]に変換して表示　
    plt.imshow(np.transpose(np_img, (1, 2, 0)))
    plt.xticks([])  # X軸の数値を消す
    plt.yticks([])  # Y軸の数値を消す
    plt.show()

# 学習用のデータローダの先頭を取り出す
data_iter = iter(train_loader)
images, _ = next(data_iter) # 返り値の２番目はラベルが入っているが描画に使わないのでダミー変数を割り当て

# 画像の表示
imshow(utils.make_grid(images))

# CNNモデル定義
class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(64 * 16 * 16, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# ハイパーパラメータの設定
learning_rate = 0.001

# モデルと最適化アルゴリズムの選択
model = SimpleCNN(num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# ハイパーパラメータの設定
num_epochs = 15

# プロット用のデータ記録配列
train_loss_history = []
train_acc_history = []
val_loss_history = []
val_acc_history = []

for epoch in range(num_epochs):
    # モデルを学習モードに変更
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0


    # ミニバッチ毎に学習を繰り返す
    for i, data in enumerate(train_loader, 0):
        # 学習データと正解ラベルを分割する
        inputs, labels = data

        # 勾配の初期化
        optimizer.zero_grad()

　　　　　# 1. モデルでの画像分類
        outputs = model(inputs)
        # 2. 正解ラベルと比較して損失を計算
        loss = criterion(outputs, labels)
        # 3. 勾配計算
        loss.backward()
        # 4. 最適化計算
        optimizer.step()

        # ミニバッチ毎で損失を累積
        running_loss += loss.item()

        # ミニバッチ毎の精度を算出、累積
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    # 学習データ全体の損失、精度を算出
    train_loss = running_loss / (i + 1)
    train_acc = correct / total

    # 検証データに対する損失と精度
    # モデルを評価モードに変更
    model.eval()
    correct = 0
    total = 0
    running_loss = 0.0

    # 勾配計算をしないよう設定
    with torch.no_grad():
        for i, data in enumerate(valid_loader, 0):
            inputs, labels = data
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item()

            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        val_loss = running_loss / (i + 1)
        val_acc = correct / total

    # 学習、検証時の損失、精度を保存
    train_loss_history.append(train_loss)
    train_acc_history.append(train_acc)
    val_loss_history.append(val_loss)
    val_acc_history.append(val_acc)
    print('Epoch {}: Train Loss: {:.4f}, Train Acc: {:.4f}, Val Loss: {:.4f}, Val Acc: {:.4f}'.format(epoch + 1,
                                                                                                      train_loss,
                                                                                                      train_acc,
                                                                                                      val_loss,
                                                                                                      val_acc))

epochs = range(1, num_epochs+1)
plt.plot(epochs, train_loss_history, lw=4, label="Training loss")
plt.plot(epochs, val_loss_history, lw=4, label="Validation loss")
plt.legend()
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.show()

plt.plot(epochs, train_acc_history, lw=4, label="Training accuracy")
plt.plot(epochs, val_acc_history, lw=4, label="Validation accuracy")
plt.legend()
plt.xlabel("Epochs")
plt.ylabel("Accuracy")
plt.show()