Pytorch – 事前学習モデルを使ってクラス分類モデルを学習する方法

目次

概要

Pytorch で事前学習済みモデルを使ってクラス分類モデルを学習する方法について解説します。

事前学習済みモデル

昨今の CNN モデルは数千万~数億のパラメータで構成されるため、このモデルのパラメータを1から調整するには、大規模なデータセットと膨大な計算リソースが要求されます。そのため、用意したデータセットのサンプルが少ない場合や潤沢な計算リソースを利用できない場合は、精度のよいモデルを作成することができません。この問題を解決するテクニックとして、転移学習 (Transfer Learning) があります。転移学習では、事前に大規模なデータセットで学習したモデルを使い、用意したデータセットでその重みを調整します。これにより、小規模なデータセットでも精度のよいモデルを手早く作成することができるようになります。

torchvision では、ImageNet で事前学習済みのモデルが提供されています。使用できるモデルの一覧は以下の記事を参照してください。

Pytorch – 学習済みモデルで画像分類を行う方法 – pystyle

転移学習の方法

基本的な CNN の構造は、画像から特徴量を抽出するための特徴抽出器 (extractor と抽出された特徴量を元に分類を行う分類器 (classifier) の2つからなります。転移学習の場合、分類器の全結合層をデータセットのクラス数に合わせて変更し、特徴抽出器は事前学習済みモデルの重みで初期化します。

学習方法については、特徴抽出器と分類器の両方のパラメータを調整する方法 (Finetuning) と特徴抽出器のパラメータは固定して、分類器のパラメータのみ調整する2通りの方法があります。

それぞれ次のような特徴があります。

  • 特徴抽出器と分類器の両方のパラメータを調整する (Finetuning)
    • データセットは中規模以上
    • 計算リソースがかかる
    • データセットの規模がそれなりにあるなら、精度は分類器のパラメータのみ調整する場合より上がりやすい
  • 分類器のパラメータのみ調整する
    • データセットは小規模でも可
    • 計算リソースが少ない

実装方法

公式チュートリアル「Transfer Learning for Computer Vision Tutorial」を参考にして進めていきます。

必要なモジュールを import する

In [1]:
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from IPython import display
from PIL import Image
In [2]:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

データセットを用意する

題材として、蟻と鉢の2クラスの画像で構成される hymenoptera データセットを使用します。 データセットは こちら からダウンロードできます。

In [3]:
# ダウンロードして、解凍する。
url = "https://download.pytorch.org/tutorial/hymenoptera_data.zip"
save_dir = Path("/data")  # 適宜変更してください
dataset_dir = save_dir / "hymenoptera_data"

datasets.utils.download_and_extract_archive(url, save_dir)

自作のデータセットを使う場合は、クラスごとに画像がディレクトリで分けられている以下のディレクトリ構造を用意してください。

my_dataset
├── train
│   ├── A: 学習用のクラス A の画像があるディレクトリ
│   ├── B: 学習用のクラス B の画像があるディレクトリ
│   └── C: 学習用のクラス C の画像があるディレクトリ
└── val
    ├── A: テスト用のクラス A の画像があるディレクトリ
    └── B: テスト用のクラス B の画像があるディレクトリ
    └── C: テスト用のクラス C の画像があるディレクトリ

用意したら、上記 [3] の部分を代わりに以下に変更してください。

dataset_dir = Path("my_dataset")  # my_dataset ディレクトリのパス

Transform を作成する

ImageNet の事前学習済みモデルで学習または推論を行う際に以下の前処理が必要です。

  • 入力の大きさを (224, 224) にする
  • 入力を RGB チャンネルごとに平均 (0.485, 0.456, 0.406)、分散 (0.229, 0.224, 0.225) で標準化する

また、学習時はランダムな切り抜き、左右反転によるデータオーグメンテーションを行います。

  • 学習時
    1. ランダムに大きさ (224, 224) で切り抜く
    2. ランダムに左右反転を行う
    3. PIL Image をテンソルにする
    4. RGB チャンネルごとに平均 (0.485, 0.456, 0.406)、分散 (0.229, 0.224, 0.225) で標準化する
  • 推論時
    1. 大きさ (256, 256) にリサイズする
    2. 大きさ (224, 224) で画像の中心を切り抜く
    3. PIL Image をテンソルにする
    4. RGB チャンネルごとに平均 (0.485, 0.456, 0.406)、分散 (0.229, 0.224, 0.225) で標準化する

これらの処理を行う Transform を作成します。

In [4]:
data_transforms = {
    # 学習時の Transform
    "train": transforms.Compose(
        [
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    ),
    # 推論時の Transform
    "val": transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    ),
}

Dataset を作成する

hymenoptera データセットを読み込む Dataset を作成します。 先ほどダウンロードして解凍したデータセットは以下のディレクトリ構造になっています。

/data/hymenoptera_data
├── train
│   ├── ants
│   └── bees
└── val
    ├── ants
    └── bees

ディレクトリがこのような構造になっている場合、ImageFolder が利用できます。

Pytorch – 自作のデータセットを扱う Dataset クラスを作る方法 – pystyle

In [5]:
# Dataset を作成する。
img_datasets = {
    x: datasets.ImageFolder(dataset_dir / x, data_transforms[x])
    for x in ["train", "val"]
}

ImageFolder はフォルダ名からクラス名及びクラス ID を作成します。 クラス名の一覧は ImageFolder.classes で取得できます。

In [6]:
class_names = img_datasets["train"].classes
print(class_names)
['ants', 'bees']

DataLoader を作成する

Pytorch – Transforms、Dataset、DataLoader について解説 – pystyle

In [7]:
dataloaders = {
    x: data.DataLoader(img_datasets[x], batch_size=4, shuffle=True, num_workers=4)
    for x in ["train", "val"]
}

デバイスを作成する

Pytorch – 計算を行うデバイスを指定する方法について – pystyle

In [8]:
def get_device(gpu_id=-1):
    if gpu_id >= 0 and torch.cuda.is_available():
        return torch.device("cuda", gpu_id)
    else:
        return torch.device("cpu")


device = get_device(gpu_id=0)

学習用のヘルパー関数を作成する

指定したエポック数だけ学習を行う train() 関数とその内部で呼び出す1エポックだけ学習する train_on_epoch() を作成します。

Pytorch – Fashion-MNIST で CNN モデルによる画像分類を行う – pystyle

In [9]:
def train(model, criterion, optimizer, scheduler, dataloaders, device, n_epochs):
    """指定したエポック数だけ学習する。
    """
    history = []
    for epoch in range(n_epochs):
        info = train_on_epoch(
            model, criterion, optimizer, scheduler, dataloaders, device
        )
        info["epoch"] = epoch + 1
        history.append(info)

        print(
            f"epoch {info['epoch']:<2} "
            f"[train] loss: {info['train_loss']:.6f}, accuracy: {info['train_accuracy']:.0%} "
            f"[test] loss: {info['val_loss']:.6f}, accuracy: {info['val_accuracy']:.0%}"
        )
    history = pd.DataFrame(history)

    return history
In [10]:
def train_on_epoch(model, criterion, optimizer, scheduler, dataloaders, device):
    """1エポックだけ学習する学習する。
    """
    info = {}
    for phase in ["train", "val"]:
        if phase == "train":
            model.train()  # モデルを学習モードに設定する。
        else:
            model.eval()  # モデルを推論モードに設定する。

        total_loss = 0
        total_correct = 0
        for inputs, labels in dataloaders[phase]:
            # データ及びラベルを計算を実行するデバイスに転送する。
            inputs, labels = inputs.to(device), labels.to(device)

            # 学習時は勾配を計算するため、set_grad_enabled(True) で中間層の出力を記録するように設定する。
            with torch.set_grad_enabled(phase == "train"):
                # 順伝搬を行う。
                outputs = model(inputs)
                # 確率の最も高いクラスを予測ラベルとする。
                preds = outputs.argmax(dim=1)

                # 損失関数の値を計算する。
                loss = criterion(outputs, labels)

                if phase == "train":
                    # 逆伝搬を行う。
                    optimizer.zero_grad()
                    loss.backward()

                    # パラメータを更新する。
                    optimizer.step()

            # この反復の損失及び正答数を加算する。
            total_loss += float(loss)
            total_correct += int((preds == labels).sum())

        if phase == "train":
            # 学習率を調整する。
            scheduler.step()

        # 損失関数の値の平均及び精度を計算する。
        info[f"{phase}_loss"] = total_loss / len(dataloaders[phase].dataset)
        info[f"{phase}_accuracy"] = total_correct / len(dataloaders[phase].dataset)

    return info

損失関数と精度の履歴を描画するヘルパー関数を作成する

損失関数と精度の履歴を描画するヘルパー関数を作成します。

In [11]:
def plot_history(history):
    fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(8, 3))

    # 損失の推移
    ax1.set_title("Loss")
    ax1.plot(history["epoch"], history["train_loss"], label="train")
    ax1.plot(history["epoch"], history["val_loss"], label="val")
    ax1.set_xlabel("Epoch")
    ax1.legend()

    # 精度の推移
    ax2.set_title("Accuracy")
    ax2.plot(history["epoch"], history["train_accuracy"], label="train")
    ax2.plot(history["epoch"], history["val_accuracy"], label="val")
    ax2.set_xlabel("Epoch")
    ax2.legend()

    plt.show()

Finetuning

モデルを作成する

今回は小規模なデータセットなので、ResNet-18 を利用します。 まず、モデルの全パラメータを学習する場合を記載します。

In [12]:
# ResNet-18 を作成する。
model_ft = models.resnet18(pretrained=True)

# 出力層の出力数を ImageNet の 1000 からこのデータセットのクラス数である 2 に置き換える。
model_ft.fc = nn.Linear(model_ft.fc.in_features, len(class_names))

# モデルを計算するデバイスに転送する。
model_ft = model_ft.to(device)

# 損失関数を作成する。
criterion = nn.CrossEntropyLoss()

# 最適化手法を選択する。
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
Downloading: “https://download.pytorch.org/models/resnet18-5c106cde.pth” to /root/.cache/torch/checkpoints/resnet18-5c106cde.pth

学習する

In [13]:
n_epochs = 10  # エポック数
history = train(
    model_ft, criterion, optimizer, scheduler, dataloaders, device, n_epochs
)

plot_history(history)
epoch 1  [train] loss: 0.129269, accuracy: 72% [test] loss: 0.074556, accuracy: 90%
epoch 2  [train] loss: 0.112779, accuracy: 82% [test] loss: 0.065001, accuracy: 91%
epoch 3  [train] loss: 0.110108, accuracy: 84% [test] loss: 0.110336, accuracy: 85%
epoch 4  [train] loss: 0.118470, accuracy: 83% [test] loss: 0.104664, accuracy: 82%
epoch 5  [train] loss: 0.089549, accuracy: 82% [test] loss: 0.090263, accuracy: 85%
epoch 6  [train] loss: 0.091574, accuracy: 84% [test] loss: 0.073524, accuracy: 90%
epoch 7  [train] loss: 0.127873, accuracy: 77% [test] loss: 0.223082, accuracy: 76%
epoch 8  [train] loss: 0.081925, accuracy: 86% [test] loss: 0.072420, accuracy: 88%
epoch 9  [train] loss: 0.094426, accuracy: 84% [test] loss: 0.059911, accuracy: 91%
epoch 10 [train] loss: 0.061152, accuracy: 89% [test] loss: 0.059724, accuracy: 90%

推論結果を表示する

いくつかのサンプル画像に対して推論を行い、結果を表示します。

In [14]:
def show_prediction(model, transform, imgs, n=3):
    for img in imgs:
        # 1. PIL Image を標準化したテンソルにする。
        # 2. バッチ次元を追加する。 (C, H, W) -> (1, C, H, W)
        # 3. 計算するデバイスに転送する。
        inputs = transform(img).unsqueeze(dim=0).to(device)

        with torch.no_grad():
            # 順伝搬を行う。
            outputs = model(inputs)

            # 確率の最も高いクラスを予測ラベルとする。
            class_id = int(outputs.argmax(dim=1)[0])

        # 推論結果を表示する。
        display.display(img.resize((224, 224)))
        print(class_names[class_id])


imgs = []
for class_dir in (dataset_dir / "val").iterdir():
    for img_path in sorted(class_dir.iterdir())[:2]:
        img = Image.open(img_path)
        imgs.append(img)

show_prediction(model_ft, data_transforms["val"], imgs)
bees
ants
bees
bees

分類器のパラメータのみ調整する

モデルを作成する

次に分類器のパラメータのみ調整する場合のコードを記載します。 特徴抽出器のパラメータは調整しないので、各パラメータの勾配を計算するかどうかを決める属性 requires_grad をすべて False に設定します。 それ以外の部分は Finetuning の場合と同様です。

In [15]:
# ResNet-18 を作成する。
model_fixed = models.resnet18(pretrained=True)
for param in model_fixed.parameters():
    param.requires_grad = False  # 勾配を計算しない

# 出力層の出力数を ImageNet の 1000 からこのデータセットのクラス数である 2 に置き換える。
model_fixed.fc = nn.Linear(model_fixed.fc.in_features, len(class_names))

# モデルを計算するデバイスに転送する。
model_fixed = model_fixed.to(device)

# 損失関数を作成する。
criterion = nn.CrossEntropyLoss()

# 最適化手法を選択する。
optimizer = optim.SGD(model_fixed.parameters(), lr=0.001, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

学習する

In [16]:
n_epochs = 10  # エポック数
history = train(
    model_fixed, criterion, optimizer, scheduler, dataloaders, device, n_epochs
)

plot_history(history)
epoch 1  [train] loss: 0.158743, accuracy: 69% [test] loss: 0.050216, accuracy: 94%
epoch 2  [train] loss: 0.125644, accuracy: 77% [test] loss: 0.140428, accuracy: 76%
epoch 3  [train] loss: 0.154087, accuracy: 75% [test] loss: 0.044722, accuracy: 93%
epoch 4  [train] loss: 0.149899, accuracy: 76% [test] loss: 0.047211, accuracy: 93%
epoch 5  [train] loss: 0.119348, accuracy: 79% [test] loss: 0.043895, accuracy: 95%
epoch 6  [train] loss: 0.144165, accuracy: 75% [test] loss: 0.074479, accuracy: 91%
epoch 7  [train] loss: 0.122264, accuracy: 78% [test] loss: 0.060641, accuracy: 92%
epoch 8  [train] loss: 0.105848, accuracy: 81% [test] loss: 0.049391, accuracy: 94%
epoch 9  [train] loss: 0.106663, accuracy: 85% [test] loss: 0.048351, accuracy: 93%
epoch 10 [train] loss: 0.104302, accuracy: 85% [test] loss: 0.059005, accuracy: 94%

コメント

コメントする

目次