MSBOOKS

プログラミングとか深層学習とか日常の出来事とか

深層学習でラブライブサンシャインのキャラクター顔認識をしてみた part2(ネットワークモデルの作成と学習)

はじめに

前回は学習データの作成をやりましたので,今回は深層学習ネットワークモデルの作成と学習をやっていきたいと思います.

ネットワークモデルの作成

今回PyTorchを使ったネットワークモデルを作成していくのですが,今回作成した学習データからラブライブサンシャインのキャラクター9人を認識させるには以下の条件を満たす必要があります.

  • 入力画像は64×64の大きさでチャンネル数は3(カラー)
  • 出力は0~8の値(0~8である各確率)

今回やりたい認識を実現するためには,Deep learningclassification(分類)のネットワークを使うとよさげ.今回は大きなスケールの画像でも高精度な結果を実現しているVGGと呼ばれるネットワーク構成を用います.以下にVGGの紹介のURLを載せておきます.
https://www.techleer.com/articles/305-vgg-16-an-advanced-approach-towards-accurate-large-scale-image-recognition/www.techleer.com

簡単に説明すると,上の画像のように入力画像に対し,Convolutionを通って特徴を抽出しながら,MaxPoolを通って画像を徐々に縮小していきます.最後にその入力画像に該当するラベル(正解)の確率を出力します.
構造はVGGを基本にしましたが,今回は入力が64×64と比較的小さいことと,最後の出力は0~8の9チャンネルしかないので,小さめのネットワーク構成にしました.

ソースコード(Model)

・Model.py

import torch
import torch.nn as nn

#ネットワークの構成
class VGG(nn.Module):
    def __init__(self, num_classes):
        super().__init__()

        self.block1_output = nn.Sequential (
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
 
        self.block2_output = nn.Sequential (
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
 
        self.block3_output = nn.Sequential (
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.block4_output = nn.Sequential (
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
 
        self.classifier = nn.Sequential(
            nn.Linear(8192, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes)
        )
 
    def forward(self, x):
        y = self.block1_output(x)
        y = self.block2_output(y)
        y = self.block3_output(y)
        y = self.block4_output(y)
        y = y.view(y.size(0), -1)
        y = self.classifier(y)
        return y

ネットワークを学習させてみる

それではPyTorchで学習をさせていきます.今回集めた900枚のデータを,学習用とテスト用にsklearnを使って分けます.それをTensorに変えてからDataLoaderに読み込ませ,バッチサイズ(今回は45)ごとにネットワークに学習させていきます.今回作成したプログラムはtrain(学習)test(テスト)のメソッドから成るクラスにしました.それぞれの処理は以下のようになっています.
●trainメソッド

  • DataLoaderから読みだす
  • 入力に対するネットワークの出力結果を得る
  • 出力結果と正解から誤差を求める
  • 誤差逆伝搬させる
  • optimizerのstepを進める

●testメソッド

  • DataLoaderから読みだす
  • 入力に対するネットワークの出力結果を得る
  • 出力結果のうち確率が最大のラベル(0~8)を求める
  • 正解であればcorrectを1増やす

ソースコード(Trainer)

・Trainer.py

import os
import numpy as np
import cv2
from tqdm  import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from Model import VGG
from sklearn.model_selection import train_test_split

OUTPUT_PATH = "output/model.pth"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #GPUを使用

class Trainer():
    def __init__(self):
        self.model = VGG(9).to(device) #モデルの定義
        self.criterion = nn.CrossEntropyLoss().to(device)
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.0001)

        #学習画像と正解を学習用とテスト用に分ける
        train_image, test_image, train_class, test_class = train_test_split(np.load('anime_face_image.npy'), np.load('anime_face_class.npy'), test_size=0.1)
        #data_loaderを作成
        train_data = torch.utils.data.TensorDataset(torch.Tensor(train_image), torch.Tensor(train_class))
        self.train_loader = torch.utils.data.DataLoader(train_data, batch_size=45, shuffle=True)
        test_data = torch.utils.data.TensorDataset(torch.Tensor(test_image), torch.Tensor(test_class))
        self.test_loader = torch.utils.data.DataLoader(test_data, batch_size=45, shuffle=True)

    def train(self,epoch):
        self.model.train()
        for batch_idx, (image, label) in enumerate(self.train_loader,1):
            label = label.type(torch.LongTensor)
            image, label = image.to(device), label.to(device)
            self.optimizer.zero_grad()
            output = self.model(image)
            loss = self.criterion(output, label)
            loss.backward()
            self.optimizer.step()

            tqdm.write('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(image), len(self.train_loader.dataset),
                100. * batch_idx / len(self.train_loader), loss.item()))

    def test(self):
        self.model.eval()
        test_loss = 0
        correct = 0
        for (image, label) in self.test_loader:
            label = label.type(torch.LongTensor)
            image, label = image.to(device), label.to(device)
            output = self.model(image)
            test_loss += self.criterion(output, label)
            _, pred = torch.max(output.data,1)
            correct += (pred == label).sum().item()

        test_loss /= len(self.test_loader.dataset)
        tqdm.write('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(self.test_loader.dataset),
            100. * correct / len(self.test_loader.dataset)))

    def saveModel(self):
        torch.save(self.model.state_dict(), OUTPUT_PATH)

if __name__ == '__main__':
    trainer = Trainer()
    try:
        os.makedirs("output/")
    except:
        pass

    for epoch in tqdm(range(5)):
        trainer.train(epoch+1)
        trainer.test()

    trainer.saveModel()

最後にoutputに学習結果を書き込むメソッドを呼び出していますが,これは単純に例ではepoch5のネットワークの状態を書き出していることになります.この書き出しを行うことで,次にやる自分でもってきた新たなテスト画像を評価できるようになります.
このネットワークで学習させた結果,今回の場合正解率92%を得ることができました.けっこうよさげ.

おわりに

今回ネットワークの構成から実際にネットワークに学習させ,未知の画像に適用した場合の評価まで実装しました.次回は自分で新たにもってきた画像の適用をするのをやっていきたいと思います.