🍕

Pytorchでピザ判定機を作る🍕

2023/04/29に公開

はじめに

こんにちは、あんどうです。今回はPytorchでピザの画像を分類する簡単な画像分類モデルを作成していきたいと思います。

ライブラリ類の準備

from torch import optim
from torch import nn # ネットワークや各種レイヤー
from torch.nn import functional #より詳しいレイヤー
from torch.utils.data import DataLoader
from torchvision import datasets # 画像データセットのモジュール
from torchvision import transforms # 画像をTorchのテンソルに変換する
from torchviz import make_dot
from torchvision import models
from torch.optim import Adam, lr_scheduler
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from datetime import timedelta
import torch
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import opendatasets as op
import time
import shutil

データセットの準備

今回使うデータセットをダウンロードします。今回はKaggleで公開されている、「Pizza or Not Pizza?」というデータセットを使います。opendatasetsというライブラリを用いてダウンロードします。 op.download()でダウンロードができますが、kaggleからのダウンロードの場合にはkaggle apiで使われるkaggle.jsonがカレントディレクトリに無いとエラーが出るので注意が必要です。

shutil.copyfile("/content/drive/MyDrive/kaggle/kaggle.json", "/content/kaggle.json") #kaggle.jsonをコピーしてカレントディレクトリに移動
op.download("https://www.kaggle.com/datasets/carlosrunner/pizza-not-pizza")

ディレクトリ構造は以下のようになります。not_pizzaにピザでない食べ物の画像、pizzaに食べ物の画像が入っています。

└── pizza_not_pizza
    ├── food101_subset.py
    ├── not_pizza [983 entries exceeds filelimit, not opening dir]
    └── pizza [983 entries exceeds filelimit, not opening dir]

Datast&DataLoaderの作成

Pytorchで学習を行うためには、DatasetとDataloaderの作成が必要になります。 データがダウンロードできたら、まずDatasetを作成していきます。最初にtorchのモジュールであるdatasetsのImageFolderメソッドを用いてデータを読み込みます。 ディレクトリがtrainデータとtestデータで別れている構造の場合、読み込んだままDatasetにすることができますが、今回は別れていないため自分で分ける必要があります。そのためにDatasetのクラスを改造してサブクラスを定義します。(インデックスでデータを分けれるようにしています。)

data_path = "/content/pizza-not-pizza/pizza_not_pizza"
# os.rmdir(data_path + "/.ipynb_checkpoints")

data = datasets.ImageFolder("/content/pizza-not-pizza/pizza_not_pizza")
#Datasetのサブクラス(trainとvalidとtestに分けるために必要)
class MySubset(torch.utils.data.Dataset):
    def __init__(self, dataset, indices, transform=None):
        self.dataset = dataset
        self.indices = indices
        self.transform = transform

    def __getitem__(self, idx):
        img, label = self.dataset[self.indices[idx]]
        if self.transform:
            img = self.transform(img)

        return img, label

    def __len__(self):
        return len(self.indices)

サブクラスを定義できたら次は画像を読み込んだ時に行う変換をtransforms.Composeの中に定義します。基本的にはテンソルに変換するだけですが、今回訓練データでは、正規化を始め、Data augumentationというモデルに頑健性を持たせるための変換(例えば画像を反転させたり、ランダムに一部を黒く塗りつぶしたり)を定義しています。また今回使うVGG_19bnという事前学習済みモデルが画像サイズ244*244で学習されたモデルであるので、そのサイズに合わせるようにリサイズする変換も盛り込んであります。

#訓練データ用
#正規化、反転、RandomErasing←ランダムで一部を消すやつ
train_transform = transforms.Compose([
                                      transforms.RandomResizedCrop(224),
                                      transforms.RandomHorizontalFlip(),
                                      transforms.ToTensor(),
                                      transforms.Normalize(0.5, 0.5),
                                      transforms.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 0.3), value=0, inplace=False)                        
])

#検証データ用
#正規化だけ
val_transform = transforms.Compose([
                                     transforms.Resize(256),
                                     transforms.CenterCrop(224),
                                     transforms.ToTensor(),
                                     transforms.Normalize(0.5, 0.5)
])

transformsが定義できたら実際にDataloaderを作っていきます。data→Dataset→Dataloaderの順で定義していきます。この時に、データを訓練データ、検証データ、テストデータに分けます。

batch_size = 32

indices = np.arange(len(data))
train_idx, tmp_idx = train_test_split(indices, test_size=0.2, random_state=2000) # 訓練データ:(検証データ+テストデータ)=8:2
valid_idx, test_idx = train_test_split(tmp_idx, test_size=0.5, random_state=2000)# 検証データ:テストデータ=5:5

train_dataset = MySubset(data, train_idx, train_transform)
val_dataset = MySubset(data, valid_idx, val_transform)
test_dataset = MySubset(data, valid_idx, val_transform)

train_loader = DataLoader(train_dataset, batch_size, shuffle=True, drop_last=True)
val_loader = DataLoader(val_dataset, batch_size, drop_last=True)
test_loader = DataLoader(test_dataset, batch_size, drop_last=True)

ちゃんとデータが分けられているかどうかを確認するために、各データセットの大きさを確認します。

print(f"train length: {len(train_dataset)}")
print(f"valid length: {len(val_dataset)}")
print(f"test length: {len(test_dataset)}")
train length: 1572
valid length: 197
test length: 197

今一度データのクラスを確認しておきます。

data.class_to_idx
{'not_pizza': 0, 'pizza': 1}

試しに画像を見てみましょう。

plt.imshow(data[1000][0])
plt.title(data.classes[data[1000][1]])
plt.show()

plt.imshow(data[1][0])
plt.title(data.classes[data[1][1]])
plt.show()

Lets学習!!

次に実際に学習をしていきます。事前学習済みモデルを読み込んだ後に、最終層を今回分類したい数に合わせます。今回は2クラス分類なので最終層のノードの数を2に設定します。(2クラス分類なので1つでもいいかも?) また学習させる準備として、モデルをGPUに乗せる、損失関数の定義、最適化手法の定義、学習率を変化させるスケジューラーの設定を行います。

lr = 0.001
momentum = 0.9


net = models.vgg19_bn(pretrained=True)

torch.manual_seed(42) # シード値の設定
torch.cuda.manual_seed(42) # シード値の設定

#最終ノードの出力を2にする
in_features = net.classifier[6].in_features
net.classifier[6] = nn.Linear(in_features, 2)

net.avgpool = nn.Identity()

device = "cuda:0" if torch.cuda.is_available() else "cpu"

net = net.to(device)

loss_CE = nn.CrossEntropyLoss()

optimizer = optim.SGD(net.parameters(), lr=lr, momentum=momentum)

scheduler = lr_scheduler.LinearLR(optimizer)

ようやく学習をします。学習の流れとしては、予測→損失を出す→勾配を求め逆伝播する→パラメータ更新→検証データで精度を確認のサイクルをぐるぐる回す感じです。

n_epochs = 20

loss_list = []
acc_train_list = []
acc_valid_list = []
acc_list = []
for epoch in range(n_epochs):
    start_time = time.time()
    net.train()
    loss_train = 0

    for data in train_loader:

        inputs, labels = data

        optimizer.zero_grad()

        output = net.forward(inputs.to(device))

        loss = loss_CE(output, labels.to(device))

        loss.backward()

        optimizer.step()

        loss_train += loss.tolist()

        predicted = torch.max(output, 1)[1].cpu()
        acc_train_list.append(((predicted == labels).sum() / len(labels)).item())


    scheduler.step()

    net.eval()
    loss_valid = 0

    with torch.no_grad():
        for data in val_loader:
            inputs, labels = data

            output = net.forward(inputs.to(device))

            loss = loss_CE(output, labels.to(device))

            loss_valid += loss.tolist()

            predicted = torch.max(output, 1)[1].cpu()
            acc_valid_list.append(((predicted == labels).sum() / len(labels)).item())

    delta = timedelta(seconds=time.time()-start_time)
    loss_list.append([loss_train, loss_valid])
    train_acc = (sum(acc_train_list) / len(acc_train_list))
    valid_acc = (sum(acc_valid_list) / len(acc_valid_list))
    acc_list.append([train_acc, valid_acc])
    print(epoch, f"\ttrain_loss: {loss_train:.5f}", f"\tvalid_loss: {loss_valid:.5f}", f"\ttime_delta: {delta}", f"\ttrain_acc: {train_acc:.5f}",  f"\tvalid_acc: {valid_acc:.5f}")
0 	train_loss: 26.40694 	valid_loss: 1.67561 	time_delta: 0:00:55.186875 	train_acc: 0.74171 	valid_acc: 0.90625
1 	train_loss: 15.89383 	valid_loss: 0.99066 	time_delta: 0:00:43.031918 	train_acc: 0.80548 	valid_acc: 0.92969
2 	train_loss: 13.32248 	valid_loss: 0.85355 	time_delta: 0:00:41.736625 	train_acc: 0.82972 	valid_acc: 0.93750
3 	train_loss: 10.70380 	valid_loss: 0.79423 	time_delta: 0:00:41.455123 	train_acc: 0.85061 	valid_acc: 0.93490
4 	train_loss: 8.77097 	valid_loss: 0.76774 	time_delta: 0:00:41.396329 	train_acc: 0.86556 	valid_acc: 0.94167
5 	train_loss: 9.04756 	valid_loss: 0.70780 	time_delta: 0:00:41.547104 	train_acc: 0.87628 	valid_acc: 0.94444
6 	train_loss: 9.02337 	valid_loss: 0.78081 	time_delta: 0:00:41.758102 	train_acc: 0.88274 	valid_acc: 0.94717
7 	train_loss: 7.63839 	valid_loss: 0.78893 	time_delta: 0:00:41.226905 	train_acc: 0.88975 	valid_acc: 0.94987
8 	train_loss: 6.84958 	valid_loss: 0.81775 	time_delta: 0:00:41.421087 	train_acc: 0.89590 	valid_acc: 0.95023
9 	train_loss: 6.33305 	valid_loss: 0.99078 	time_delta: 0:00:41.423647 	train_acc: 0.90128 	valid_acc: 0.95104
10 	train_loss: 5.67506 	valid_loss: 0.75765 	time_delta: 0:00:41.720734 	train_acc: 0.90625 	valid_acc: 0.95218
11 	train_loss: 5.45034 	valid_loss: 0.68942 	time_delta: 0:00:41.822234 	train_acc: 0.91029 	valid_acc: 0.95312
12 	train_loss: 5.18427 	valid_loss: 0.89161 	time_delta: 0:00:41.440836 	train_acc: 0.91415 	valid_acc: 0.95393
13 	train_loss: 5.67018 	valid_loss: 0.64947 	time_delta: 0:00:48.785079 	train_acc: 0.91723 	valid_acc: 0.95461
14 	train_loss: 4.82825 	valid_loss: 0.73357 	time_delta: 0:00:41.433255 	train_acc: 0.92007 	valid_acc: 0.95521
15 	train_loss: 3.77444 	valid_loss: 0.79571 	time_delta: 0:00:41.470007 	train_acc: 0.92335 	valid_acc: 0.95573
16 	train_loss: 3.42928 	valid_loss: 0.70412 	time_delta: 0:00:43.706830 	train_acc: 0.92617 	valid_acc: 0.95558
17 	train_loss: 3.80916 	valid_loss: 0.88906 	time_delta: 0:00:41.943897 	train_acc: 0.92847 	valid_acc: 0.95515
18 	train_loss: 3.23138 	valid_loss: 0.76953 	time_delta: 0:00:42.778128 	train_acc: 0.93085 	valid_acc: 0.95587
19 	train_loss: 4.04961 	valid_loss: 0.71628 	time_delta: 0:00:41.386932 	train_acc: 0.93275 	valid_acc: 0.95625

モデルが学習できたら、保存します。

torch.save(net.state_dict(), "/content/drive/MyDrive/Mymodel/PIzzaNet.pkl")
del net, inputs, labels, loss, output
torch.cuda.empty_cache()

学習結果を確認していきます。下の図が損失のグラフ、更に下が正解率のグラフとなっています。正解率を見る限り、エポック数をもう少し増やしても良かったかもしれません。またどちらの図も訓練データの精度のほうが検証用データの精度より低いのは、Data augumentationをしているからだと考えられます。

plt.plot(np.array(loss_list)[:, 0], label="train")
plt.plot(np.array(loss_list)[:, 1], label="valid")
plt.legend()
plt.grid()
plt.xlabel("Epoch")
plt.ylabel("Cross Entropy")
plt.title("loss")

plt.plot(np.array(acc_list)[:, 0], label="train")
plt.plot(np.array(acc_list)[:, 1], label="valid")
plt.legend()
plt.grid()
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.title("Accuracy")

結果の確認

保存したモデルを読み込んで、実際にテストデータで精度を確認していきましょう。

net = models.vgg19_bn(pretrained=False)
#最終ノードの出力を2にする
device = "cuda:0" if torch.cuda.is_available() else "cpu"
in_features = net.classifier[6].in_features
net.classifier[6] = nn.Linear(in_features, 2)
net.load_state_dict(torch.load("/content/drive/MyDrive/Mymodel/PIzzaNet.pkl"))
net.to(device)

テストデータに対して予測させます。

test_output = []
test_label = []
net.eval()

with torch.no_grad():
    for data in test_loader:
        inputs, labels = data
        output = net.forward(inputs.to(device))
        test_output.extend(output.tolist())
        test_label.extend(labels.tolist())
predictions = np.argmax(test_output, axis=1)
accuracy_score(test_label, predictions)
0.9635416666666666

正解率は約96%でした。これならピザかピザじゃないか判定できそうです。

参考文献

前処理系: https://pystyle.info/pytorch-list-of-transforms/
データセット: https://www.kaggle.com/datasets/carlosrunner/pizza-not-pizza
データセットを分けるためのコード: https://pystyle.info/pytorch-split-dataset/

Discussion