Pytorch + EfficientNetで友人2人の顔を判別する画像分類モデルを作ってみた

はじめに

※1年くらい前にやったネタを今更記事にしました。古い情報もあるかもだけど許してね

iPhoneの写真には、ピープルという機能がありまして、保存している写真に誰が写ってるのかを整理してくれます。
ある日、この機能で写真を眺めていたところ、時折友人M氏とY氏を間違えて判定しているところがありました。
2人の顔はそこまで似ていないと思ったので、実際にAIモデルで2人の顔画像を学習してみたら、うまく判定できるのか？
そう思い、Pytorchの勉強もかねて、EfficientNetをFine Tuningしてみて、精度を測ってみることにしました。

実行環境

なんとローカルで実行してます。そろそろアップグレードしたい…

OS : Win10
CPU : Intel Core i7-7700k
CPUメモリ : 16GB
GPU : GeForce GTX1070
GPUメモリ : 8GB

データセットの用意

データセットはiPhoneに入っている、友人M氏とY氏が写っている写真を、適当に手動で顔部分だけトリミングしたものを使用しました。
大きさは幅3000pxを超えるものや、100px程度のものなど、バラバラです。
これらの画像データを_input/images_のディレクトリに格納します。
M氏の画像118枚、Y氏の画像128枚用意しました。

また、画像データと目的変数をまとめたCSVファイルを作成し、_input/dataset.csv_として格納します。

id	class
m (0).jpg	M
m (1).jpg	M
m (2).jpg	M
m (3).jpg	M
m (4).jpg	M
...	...
y (123).jpg	Y
y (124).jpg	Y
y (125).jpg	Y
y (126).jpg	Y
y (127).jpg	Y

データセットと前処理

データセットクラス

Pytorchで学習を行うためには、データセットクラスを定義する必要があります。
データセットクラスはPytorchのDatasetクラスを継承するようにしています。

import cv2
import torch.utils.data as data

class MYDataset(data.Dataset):
    """
    PyTorchのDatasetクラスを継承。

    Attributes
    ----------
    file_path_df : pd.DataFrame
        読み込む画像データの情報を格納しているデータフレーム
    transform : object
        前処理クラスのインスタンス
    phase : 'train' or 'test'
        訓練か評価かを設定する。
    """

    def __init__(self, file_path_df, root_path="input/", transform=None, phase='train'):
        self.file_path_df = file_path_df  # ファイルパスが格納されているDF
        self.root_path = root_path  # ルートのパス
        self.transform = transform  # 前処理クラスのインスタンス
        self.phase = phase  # train or testの指定

    def __len__(self):
        '''画像の枚数を返す'''
        return len(self.file_path_df)

    def __getitem__(self, index):
        '''
        前処理をした画像のTensor形式のデータとラベルを取得
        '''

        # index番目の画像をロード
        img_file_name = self.file_path_df["id"][index]

        # 画像データのパス
        img_path = self.root_path + "images/" + img_file_name

        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 画像の前処理を実施
        img_transformed = self.transform(
            img, self.phase)  # torch.Size([3, 224, 224])

        # 画像のラベルを取得
        # Mだったら0、それ以外(Y)だったら1となるようにする
        label = 0 if self.file_path_df["class"][index] == "M" else 1

        return img_transformed, label

前処理クラス

上記のデータクラスの引数で使用している前処理のクラスを定義します。
前処理では、Pytorch公式のtorchvision.transformsがありますが、今回はデータ拡張の種類を増やすため、albumentationsというライブラリを使用します。

詳しい説明は省きますが、訓練時には反転・回転や、ガウスノイズ・ボカシ・色調変化などをランダムで行ってから、リサイズとトリミング(クロップ)を行うようにし、評価時はリサイズとトリミングのみを行うようにしてランダム性がないようにします。

import albumentations as A
from albumentations.pytorch import ToTensorV2

class ImageTransform():
    def __init__(self, resize, crop_size):
        self.data_transform = {
            'train': A.Compose([
                A.Transpose(p=0.5),
                A.HorizontalFlip(p=0.5),
                A.VerticalFlip(p=0.5),
                A.Rotate(limit=180, p=1.0),
                A.OneOf([
                    A.GaussNoise(p=1),
                    A.GaussianBlur(p=1),
                ], p=0.5),
                A.RandomBrightnessContrast(
                    brightness_limit=(-0.2, 0.2), 
                    contrast_limit=(-0.2, 0.2), 
                    p=0.5),
                A.HueSaturationValue(
                    hue_shift_limit=5, 
                    val_shift_limit=5, 
                    p=0.5),
                A.SmallestMaxSize(resize),
                A.RandomCrop(crop_size, crop_size),
                A.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                    max_pixel_value=255.0,
                    p=1.0),
                ToTensorV2(p=1.0),
            ], p=1.0),
            'test': A.Compose([
                A.SmallestMaxSize(resize),
                A.CenterCrop(crop_size, crop_size),
                A.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                    max_pixel_value=255.0,
                    p=1.0),
                ToTensorV2(p=1.0),
            ], p=1.0),
        }

    def __call__(self, img, phase='train'):
        """
        Parameters
        ----------
        phase : 'train' or 'test'
            前処理のモードを指定。
        """
        return self.data_transform[phase](image=img)["image"]

友人M氏、Y氏の顔画像を載せると怒られそうなので、代わりに僕が書いたアイコンの画像で前処理を実行してみます。

元画像

前処理後（訓練時）

前処理後（評価時）

イラストだとわかりづらいけど、こんな感じ。

データローダーの作成

上記のデータセットクラス・前処理クラスを_utils/dataloader_image_classification.py_に記載しておき、これを使ってデータローダーのインスタンスを作成します。
訓練用と評価用を8:2の割合で分割して作ります。

# パッケージのインポート
import pandas as pd
from sklearn.model_selection import train_test_split
from utils.dataloader_image_classification import ImageTransform, MYDataset

# CSVファイルの読み込み
df = pd.read_csv("input/dataset.csv")

# 訓練用と評価用で分割
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

# indexの初期化
train_df = train_df.reset_index()[["id", "class"]]
test_df = test_df.reset_index()[["id", "class"]]

# Datasetを作成する
crop_size = 224
resize = 224
train_dataset = MYDataset(
    file_path_df=train_df, transform=ImageTransform(resize, crop_size), phase='train')
test_dataset = MYDataset(
    file_path_df=test_df, transform=ImageTransform(resize, crop_size), phase='test')

# DataLoaderを作成する
batch_size = 32
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)

# 辞書オブジェクトにまとめる
dataloaders_dict = {"train": train_dataloader, "test": test_dataloader}

学習モデルの作成

学習を行うためのモデルを定義します。
Pytorchのほかに、EfficientNetの事前学習済みモデルを取得するために、timmというライブラリを使用します。

# パッケージのインポート
import timm
import torch
import torch.nn as nn
import torch.optim as optim

# モデルのインスタンスを生成
net = timm.create_model('efficientnet_b0', pretrained=True)

# 最後の出力層の出力ユニットをクラス(M氏 or Y氏)の2つに付け替える
net.classifier = nn.Linear(in_features=1280, out_features=2, bias=True)

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

次に学習モデルで使用する損失関数、最適化関数を設定します。
元の論文では、最適化関数はRMSPropsを使用している？みたいでしたが、なんかいろいろと難しいことが書いてあったので、今回はシンプルにSGDで学習させます。

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

# 最適化手法の設定
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

# スコア格納のためのリスト作成
tr_loss_list, tr_acc_list, va_loss_list, va_acc_list = [],[],[],[]

学習の実行

学習を行うための関数を定義します。

def train_model(net, dataloaders_dict, criterion, optimizer, num_epochs, 
                model_path="output/model.pth"):
    # 初期設定
    # GPUが使えるかを確認
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("使用デバイス：", device)

    # ネットワークをGPUへ
    net.to(device)

    # ネットワークがある程度固定であれば、高速化させる
    torch.backends.cudnn.benchmark = True

    # epochのループ
    for epoch in range(num_epochs):
        print('-------------')
        print('Epoch {}/{}'.format(epoch + 1, num_epochs))

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

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

            predlist = torch.zeros(0, dtype=torch.long, device='cpu')
            lbllist = torch.zeros(0, dtype=torch.long, device='cpu')

            # データローダーからミニバッチを取り出すループ
            for inputs, labels in dataloaders_dict[phase]:
                # GPUが使えるならGPUにデータを送る
                inputs = inputs.to(device)
                labels = labels.to(device)

                # optimizerを初期化
                optimizer.zero_grad()

                # 順伝搬（forward）計算
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = net(inputs)
                    loss = criterion(outputs, labels)  # 損失を計算
                    _, preds = torch.max(outputs, 1)  # ラベルを予測

                    # 訓練時はバックプロパゲーション
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

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

                    predlist = torch.cat([predlist, preds.detach().view(-1).cpu()])
                    lbllist = torch.cat([lbllist, labels.view(-1).cpu()])

            # epochごとのlossと正解率を表示
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double().to(torch.device("cpu")) / len(dataloaders_dict[phase].dataset)

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

            # モデルを保存
            torch.save(net.state_dict(), model_path)


            # epochごとのlossと正解率を表示
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double(
            ) / len(dataloaders_dict[phase].dataset)
            
            # リストに格納
            if phase == 'train':
                tr_loss_list.append(epoch_loss)
                tr_acc_list.append(epoch_acc)
            elif phase == 'test':
                va_loss_list.append(epoch_loss)
                va_acc_list.append(epoch_acc)

            print('  {} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))
            
            # モデルを保存
            torch.save(net.state_dict(), model_path)

上記の関数を実行します。
今回はデータ数が少ないため、データオーギュメンテーションの幅を広げるために、epoch数は300と多めにしてます。
これにより、過学習を防ぎ、より汎用的なモデルを作成できることが見込まれます。

私の環境だと、実行にGPU使って1時間弱くらいかかりました。

# 学習・検証を実行する
num_epochs=300
train_model(net, dataloaders_dict, criterion, optimizer, 
            num_epochs=num_epochs, model_path="output/efficient.pth")

予測

学習したモデルを使用し、予測を行います。
まず、予測の関数を定義します。
この関数では、予測値の他に、予測確率も返すようにしておきます。

def predict(model, test_dataloader):
    # GPUが使えるかを確認
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    sm = nn.Softmax()
    pred_list=torch.zeros(0, dtype=torch.long, device='cpu')
    pred_probas_list=torch.zeros(0, dtype=torch.long, device='cpu')
    for inputs, labels in test_dataloader:
        # GPUが使えるならGPUにデータを送る
        inputs = inputs.to(device)
        with torch.set_grad_enabled(False):
            outputs = model(inputs)
            outputs = sm(outputs)
            _, preds = torch.max(outputs, 1)  # ラベルを予測
        pred_list = torch.cat([pred_list, preds.detach().view(-1).cpu()])
        pred_probas_list = torch.cat([pred_probas_list, outputs.cpu()])
        
    return np.array(pred_list), np.array(pred_probas_list)

上記の関数を実行し、正解率を算出します。

# 予測の実行
pred, proba = predict(net, test_dataloader)

# 正解率で評価
from sklearn.metrics import accuracy_score
print("Accuracy: ", accuracy_score(test_df["class"], pred))

出力結果

Accuracy: 0.9800

かなり高い精度になりました。

どうやら、評価用データセット50枚の画像のうち、1枚だけ間違えて判定しただけで、他の画像は合っていたようです。
その間違えた画像がこちら。

※顔だけの部分だけ差し替えてます。

顔の向きとかでうまく判定できなかったのかも…

まとめ

意外と適当な画像で、適当な前処理で学習させても割とうまくいくもんですね。
EfficientNetも2019年当時は最強とか言われてたみたいですが、その後もいろいろなモデルが出てるみたいで、そこらへんとかでも試してみたいところです。
余裕があったらViTとか特性の違うモデルとかで試してみたのも記事にするかも。

参考文献

書籍「つくりながら学ぶ! PyTorchによる発展ディープラーニング」（小川雄太郎、マイナビ出版 、19/07/29)
https://github.com/YutaroOgawa/pytorch_advanced