DeepLearningによる画像分類タスクの実験です。 ResNet50を使って、画像分類を行ないます。

今回は下記の2種類のタスクを行ないました。 予想通り前者は難しく学習が安定しなかったので、後者のタスクを加えました。

• イラストの著者分類
• イラストのキャラクター分類

環境はGoogle Colabを利用します。 毎回データをアップロードするのが手間だったりしますが、実験用途で使う分には、トータルで利便が勝ります。

ResNet50に入力できる画像のサイズは224×224となります。そのためイラスト全体を使用すると縦横比がおかしくなったり、細かいパーツやタッチに関する情報が潰れてしまう懸念がありました。そこで学習の対象を「顔」のみとするべく、縦横1対1の比率で切り抜きを行ないました。画像サイズは学習時の事前処理でリサイズするので、この段階で揃えません。

著者分類では、自分が描いたイラストとその他の著者が描いたイラストを約50枚ずつ用意しました。描かれているキャラクターはある程度重複するようにしています。

Google Colabへのアップロードは、ファイル単位のみみたいなので、ローカルでzip圧縮してからアップロードして、Google Colab上で展開しました。

!unzip data.zip

自分が描いたイラストはサンプルとして提供しても良いのですが、他人が描いたイラストの方は当然提供できないので、保留します。

標準では入っていないライブラリを利用するので、インストールしておきます。

pip install pytorch-gradcam

ここからコーディングとなります。まずは学習モデルを用意します。 torchvisionから構築済のResNet50モデルを利用できます。 さらにImageNetで事前学習した重みも利用可能ですが、今回は汎用課題を対象としていないので、事前学習の重みは使いません。

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision.models import ResNet50_Weights
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import os
import pickle
from sklearn.metrics import accuracy_score
 
def get_device():
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')
 
device= get_device()
# 学習済の重みは使わないので、引数でWeight指定なし
model = models.resnet50().to(device)

下記で学習モデルの構造を確認できます。 動作させる上で必須ではありませんが、後段で学習モデル内の特定の層を指定する必要があるので、ここで情報を出力しておきます。

from torchvision.models import feature_extraction
 
feature_extraction.get_graph_node_names(model)
model

学習データを扱えるように、データセットとデータローダを用意します。 データセットは生データと正解ラベルを構造化するもので、自前で用意する必要があるのですが、 torchvisionのImageFolder機能を使えば、自動的にデータセットを作成してくれます。 正解ラベル名のフォルダの配下にそのラベルと対応する画像を入れる、というルールに従って、データセットのフォルダを用意すれば、簡単にデータセットが作れます。

データローダはデータセットを指定して、ミニバッチ数やシャッフルの有無など、データの提供方法をパラメータで調整します。 併せて、画像に対する事前処理も実施します。 この事前処理にあるNormalize(正規化)の値は、今回はスクラッチ学習のため、全チャネルについて平均と標準偏差を0.5に設定します。 事前学習済モデルを用いる場合は、そのモデルが使用したデータセットに適した値を設定する必要があります。 正規化をすると色合いがおかしくなりますが、正規化しないと学習精度が割と大きく悪化しますので、設定が必要です。

# 画像に対する事前処理
# ResNet50向けに、224*224にリサイズする
preprocess = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.5, 0.5, 0.5],
        std=[0.5, 0.5, 0.5]
    )
])
 
train_image_dir = './data/train'
test_image_dir = './data/test'
 
train_dataset = torchvision.datasets.ImageFolder(root=train_image_dir, transform=preprocess)
test_dataset = torchvision.datasets.ImageFolder(root=test_image_dir, transform=preprocess)
 
batch_size = 2
train_dataLoader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True
)
test_dataLoader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False
)

学習を行ないます。 学習状況を確認するために、エポック単位でテスト用データによる推論も実施します。 学習させるデータの数や実行環境にも大きく依存しますが、この処理は時間が掛かりますのでご注意ください。

# 学習と推論に向けた準備
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001) # Instantiate an optimizer
 
# 学習と推論のエポックを回す
for epoch in range(20):
    test_acc = 0.
 
    # 学習
    for (x, t) in train_dataLoader:
        x, t = x.to(device), t.to(device)
        model.train()
        preds = model(x)
        loss = criterion(preds, t)
        # モデル内に保持している勾配情報をリセットする
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
    # 推論
    for (x, t) in test_dataLoader:
        x, t = x.to(device), t.to(device)
        model.eval()
        preds = model(x)
        test_acc += accuracy_score(t.tolist(), preds.argmax(dim=-1).tolist())
 
    # 正解率の表示
    test_acc /= len(test_dataLoader)
    print('Epoch: {}, Valid Acc: {:.3f}'.format(
        epoch+1,
        test_acc
    ))

下記のように正解率が収束すれば、学習が上手くいっている可能性が高いです。

Epoch: 1, Valid Acc: 0.500
Epoch: 2, Valid Acc: 0.500
Epoch: 3, Valid Acc: 0.500
Epoch: 4, Valid Acc: 0.500
Epoch: 5, Valid Acc: 1.000
Epoch: 6, Valid Acc: 1.000
Epoch: 7, Valid Acc: 1.000
Epoch: 8, Valid Acc: 1.000
Epoch: 9, Valid Acc: 1.000
Epoch: 10, Valid Acc: 1.000

学習したモデルを用いて、再度推論を実施し、その判断根拠をGrad-CAMで可視化します。 利用する画像は、テスト用データローダに登録しているものです。なので著者に依らない画像を利用します。

from gradcam.utils import visualize_cam
from gradcam import GradCAM, GradCAMpp
import matplotlib.pyplot as plt
import numpy as np
from torchvision.utils import make_grid
 
target_layer = model.layer2[-1].conv2
gradcam = GradCAM(model, target_layer)
 
# 表示用の画像情報を格納する配列(オリジナル、ヒートマップ、両者のオーバレイ画像)
images_to_display = []
 
# テスト用のデータローダからミニバッチ単位でループします
for (x_batch, t_batch) in test_dataLoader:
    x_batch, t_batch = x_batch.to(device), t_batch.to(device)
 
    # モデルを使って推論を行ないます
    model.eval()
    with torch.no_grad():
        outputs = model(x_batch)
    # 予測した分類結果を格納します(ミニバッチのデータサイズ分)
    predicted_classes = outputs.argmax(dim=-1)
 
    # ミニバッチ内にある画像単位でループします
    for i in range(x_batch.size(0)):
        # Grad-CAM向けにバッチ処理用の次元を追加。形状は (1, C, H, W)
        single_image_tensor = x_batch[i].unsqueeze(0)
        predicted_class_idx = predicted_classes[i].item()
 
        # grad-CAMのマスクを取得します。マスクの形状は (1, 1, H, W)
        # 引数で渡す画像のテンソルは、データローダの事前処理で正規化済なので、そのまま使います
        mask, _ = gradcam(single_image_tensor, class_idx=predicted_class_idx)
 
        # CAMを可視化します
        heatmap_tensor, result_overlay_tensor = visualize_cam(mask.squeeze(), single_image_tensor.squeeze(0))
 
        # 結果表示用の配列に追加します
        images_to_display.extend([
            single_image_tensor.squeeze(0).cpu(),  # 元画像 (3, H, W)
            heatmap_tensor.cpu(),                  # ヒートマップ (3, H, W)
            result_overlay_tensor.cpu()            # オーバレイ画像 (3, H, W)
        ])
 
# 結果を表形式で出力します
grid_image = make_grid(images_to_display, nrow=3)
transforms.ToPILImage(mode='RGB')(grid_image)

続いて、データローダを使わず、直接画像ファイルを使って推論して、grad-CAMによる可視化を行ないます。 このモチベーションの1つは、正規化する前のオリジナル画像を併記したいからです。

import glob
from gradcam.utils import visualize_cam
from gradcam import GradCAM, GradCAMpp
from torchvision.utils import make_grid
 
images_to_display = []
result_to_display = []
 
for path in glob.glob("data/test/true/*"):
    img = Image.open(path)
    torch_img = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])(img).to(device)
    normed_torch_img = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])(torch_img)[None]
 
    # モデルを使って推論を行ないます
    # これはGrad-CAMによる可視化で必須ではありませんが、推論結果も出力したいので、実施しています
    model.eval()
    with torch.no_grad():
        outputs = model(normed_torch_img)
    # 予測した分類結果を格納します
    # ResNet50の出力層を差し替えずに使っているので、出力クラス数は1000ですが、実際に使っているのは先頭2個のみです
    predicted_classes = outputs.argmax(dim=-1)
 
    # 推論結果表示用の配列に追加します
    result_to_display.append(predicted_classes.item());
 
    target_layer = model.layer1[-1].conv2
    gradcam = GradCAM(model, target_layer)
    mask, _ = gradcam(normed_torch_img)
    heatmap_tensor_1, result_overlay_tensor_1 = visualize_cam(mask, torch_img)
 
    target_layer = model.layer2[-1].conv2
    gradcam = GradCAM(model, target_layer)
    mask, _ = gradcam(normed_torch_img)
    heatmap_tensor_2, result_overlay_tensor_2 = visualize_cam(mask, torch_img)
 
    target_layer = model.layer3[-1].conv2
    gradcam = GradCAM(model, target_layer)
    mask, _ = gradcam(normed_torch_img)
    heatmap_tensor_3, result_overlay_tensor_3 = visualize_cam(mask, torch_img)
 
    target_layer = model.layer4[-1].conv2
    gradcam = GradCAM(model, target_layer)
    mask, _ = gradcam(normed_torch_img)
    heatmap_tensor_4, result_overlay_tensor_4 = visualize_cam(mask, torch_img)
 
    # Grad-CAM表示用の配列に追加します
    images_to_display.extend([
        torch_img.squeeze(0).cpu(),              # 元画像 (3, H, W)
        heatmap_tensor_1.cpu(),                  # ヒートマップ (3, H, W)
        result_overlay_tensor_1.cpu(),           # オーバレイ画像 (3, H, W)
        heatmap_tensor_2.cpu(),                  # ヒートマップ (3, H, W)
        result_overlay_tensor_2.cpu(),           # オーバレイ画像 (3, H, W)
        heatmap_tensor_3.cpu(),                  # ヒートマップ (3, H, W)
        result_overlay_tensor_3.cpu(),           # オーバレイ画像 (3, H, W)
        heatmap_tensor_4.cpu(),                  # ヒートマップ (3, H, W)
        result_overlay_tensor_4.cpu()            # オーバレイ画像 (3, H, W)
    ])
 
# 推論結果を出力します
print(result_to_display)
 
# Grad-CAMを表形式で出力します
grid_image = make_grid(images_to_display, nrow=9)
transforms.ToPILImage(mode='RGB')(grid_image)

著者分類は全く上手くいきませんでしたが、自分以外の著者のイラストについて、特定の著者に統一しなかったことが学習が安定しなかった原因かもしれません。まあ、元々難しいタスクなので、それだけが原因とは思いません。しかし、自分以外の著者のイラストの範囲が広すぎて、自分に近い画風の人とかけ離れた画風の人があった場合、上手くいかないように思えます。