ディープラーニングを使用してファッションアイテムの画像を分類する（後編）

はじめに
結果
まとめ
参考リンク

はじめに

後編では、PyTorchというフレームワーク上で畳み込みニューラルネットワーク（CNN）を使ったディープラーニングの手法を利用して、実際にファッションアイテムの画像をカテゴリ分類してみます。
用途としては、通販サイトの商品登録を楽にするために、商品画像をアップロードすると商品カテゴリを自動でサジェストする機能に使うためのモデルを想定しています。

データセット

Fashion-MNISTという、ファッションアイテムの画像が10クラスにラベル付けされているデータセットを使用します。
28×28のグレースケール画像が、訓練用に6万点、テスト用に1万点用意されています。
github.com

モデル

PyTorchには既に大規模なデータセットによって学習済みのモデルが複数実装されています。この中から「VGG-16」というモデルを、今回利用するデータセットに合わせて転移学習させます。転移学習とは、とある用途のために既に学習させられているモデルを、別の用途に使うために再学習させることです。転移学習によって、小規模なデータセットで高い精度を実現させることや、学習にかける時間を短縮させることが期待できます。

バージョン・環境

利用するPyTorchのバージョンは「1.0」です。
またJupyter Notebookで実行することを想定しています。

実装

使用するパッケージをインポートします。

from tqdm import tqdm
import matplotlib.pyplot as plt
%matplotlib inline
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision
from torchvision import models, transforms

データセットを定義します。
今回使うFashion-MNISTのデータセットは28×28のグレースケール画像ですが、VGG-16は224×224のRGB画像を入力するように作られているため、画像を加工する必要があります。

# バッチサイズの定義
batch_size=32
# 画像加工の定義
transform = transforms.Compose([transforms.Resize(224),
                                transforms.ToTensor(),
                                transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
                                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])

# データセットのロード
train_dataset = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

test_dataset = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)
test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

data_loaders_dict = {"train": train_data_loader, "test": test_data_loader}

モデルを定義します。VGG16は1000クラスの分類を行うように作られているため、Fashion-MNISTに合わせて10クラスの分類を出力するように変更します。

# 学習済みのVGG-16モデルをロード
# VGG-16モデルのインスタンスを生成
use_pretrained = True  # 学習済みのパラメータを使用
net = models.vgg16(pretrained=use_pretrained)

# 最後の出力層の出力ユニットをFashion-MNISTの10個につけかえる
net.classifier[6] = nn.Linear(in_features=4096, out_features=10)

損失関数などの設定をします。画像分類では標準的な、損失関数にはクロスエントロピー誤差、最適化手法にはSGDを使用します。

# 訓練モードに設定
net.train()

# 損失関数の設定
criterion = nn.CrossEntropyLoss()

# 最適化手法の設定
optimizer = optim.SGD(params=params_to_update, lr=0.001, momentum=0.9)

VGG-16は画像から特徴量を抽出するネットワークと特徴量を元に分類を行うネットワークに分かれていますが、このうち分類を行うネットワークのみ再学習させるよう、重み（weight）とバイアス（bias）の固定を解除します。

# 転移学習で学習させるパラメータを、変数params_to_updateに格納する
params_to_update = []

# 学習させるパラメータ名
update_param_names = ["classifier.0.weight", "classifier.0.bias", "classifier.3.weight", "classifier.3.bias", "classifier.6.weight", "classifier.6.bias"]

# 学習させるパラメータ以外は勾配計算をなくし、変化しないように設定
for name, param in net.named_parameters():
    if name in update_param_names:
        param.requires_grad = True
        params_to_update.append(param)
        print(name)
    else:
        param.requires_grad = False

# モデルを学習させる関数を作成
def train_model(net, data_loaders_dict, criterion, optimizer, num_epochs):

    net.to(device)

    # 学習結果を保存するhistoryの初期化
    history = {}
    history['train_loss_values'] = []
    history['train_accuracy_values'] = []
    history['test_loss_values'] = []
    history['test_accuracy_values'] = []

    # 学習と検証（epoch）のループ
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-------------')

        for phase in ['train', 'test']:
            if phase == 'train':
                net.train()  # モデルを訓練モードに
            else:
                net.eval()   # モデルを検証モードに

            epoch_loss = 0.0  # epochの損失和
            epoch_corrects = 0  # epochの正解数

            # データローダーからミニバッチを取り出すループ
            for inputs, labels in tqdm(data_loaders_dict[phase]):

                inputs = inputs.to(device)
                labels = labels.to(device)

                # optimizerを初期化
                optimizer.zero_grad()

                # 順伝搬（forward）計算
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = net(inputs)
                    _, preds = torch.max(outputs, 1)  # ラベルを予測
                    loss = criterion(outputs, labels)  # 損失を計算
                    
                    # 訓練時はバックプロパゲーション
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                    # イテレーション結果の計算
                    # lossの合計を更新
                    epoch_loss += loss.item() * inputs.size(0)  
                    # 正解数の合計を更新
                    epoch_corrects += torch.sum(preds == labels.data)

            # epochごとのlossと正解率を表示
            epoch_loss = epoch_loss / len(data_loaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double(
            ) / len(data_loaders_dict[phase].dataset)

            # エポック毎のlossと正解率をhistoryに格納する
            history[f'{phase}_loss_values'].append(epoch_loss)
            history[f'{phase}_accuracy_values'].append(float(epoch_acc.cpu().numpy()))

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))
    return history

学習を行う回数（エポック数）を指定し、学習を実行します。

# エポック数を定義
num_epochs=10

# 学習を実行（結果をhistoryに格納）
history = train_model(net, data_loaders_dict, criterion, optimizer, num_epochs=num_epochs)

学習結果をグラフに表示します。

plt.plot(range(num_epochs), history['train_accuracy_values'], label='train_acc')
plt.plot(range(num_epochs), history['test_accuracy_values'], label='test_acc')
plt.legend()
plt.grid()
plt.show()