Pytorchでピザ判定機を作る🍕
はじめに
こんにちは、あんどうです。今回は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