PyTorchで手書き文字(MNIST)の認識の実装に挑戦【初級 深層学習講座】
いまさら感がありますが、MINSTのデータセットを使った手書き文字認識をPyTorchを使ってゼロから実装してみました。この記事では、スクラッチからの手書き文字認識のモデルを実装にチャレンジします。加えて、TIMMを活用して実装する例も紹介し、ライブラリを利用したモデル構築の方法も解説します。
はじめに
MNISTの手書き数字のデータセットは、「0」〜「9」の数字を手書きしたデータセットです。
MNISTは、ディープラーニングの学習用によく利用されているため、手軽にダウンロードしてデータセットとして利用する環境が整っています。
ちょっとした実験をする際にMNISTの手書き文字データセットをよく利用します。学習時間が短い割には、データ数が7万枚と多いため、実験に利用しやすいからです。このデータセット、覚えておくと重宝します。
今回は、PyTorchで手書き文字を識別するモデルを作成し、学習させてみたいと思います。
普段は、TIMM(PyTorch Image Models)などを使ってベースモデルを加工してモデルを作成しますが、今回はスクラッチからモデルを作成したいと思います。
コードはGoogle Colab or Jupyter notebookで動かす前提としています。
Google Colabで動作するコードはこちらに置いています
ライブラリのインストール
この記事で利用するライブラリをインポートします。インポートするのは、torch関連、torchvision, matplotlib, numpy, tqdm, sklearnです。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import matplotlib.pyplot as plt
import numpy as np
from tqdm.notebook import tqdm
from sklearn.metrics import accuracy_score
Google Colab/Jupyterで動かさない場合は、tqdm.notebook
をtqdm
に変えてください。コードは以下になります。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
from sklearn.metrics import accuracy_score
importでエラーが出た場合は、必要なライブラリをインストールしてください
MNIST(手書き数字データセット)
データセット概要
MNISTのデータセットは以下の構成です
- 訓練データ(6万枚)、テストデータ(1万枚)の合計7万枚
- 8bitグレースケール(0~255)、28×28画素
PyTorchの場合は、torchvision.datasets.MNIST()
を使うことで簡単にダウンロードし、データセットとして利用するこtが可能です。
データセットのダウンロード
以下のコードでtrain_dataset
に6万枚の訓練データが、valid_dataset
に1万枚のテストデータがダウンロードされます。
ローカル環境の場合、初回はデータセットがダウンロードされるので、インターネットへの接続が必要です
train_dataset = torchvision.datasets.MNIST(root="data",
train=True,
transform=torchvision.transforms.ToTensor(),
download=True)
valid_dataset = torchvision.datasets.MNIST(root="data",
train=False,
transform=torchvision.transforms.ToTensor(),
download=True)
データの中身を確認してみます。以下のコードを実行すると、訓練データの25枚の内容が表示されます。
# データを確認
fig, ax = plt.subplots(5,5, figsize=(10,10))
for i in range(25) :
img, label = train_dataset[i]
r, c = i//5, i%5
ax[r, c].imshow(img.squeeze(), cmap="gray")
ax[r, c].axis("off")
ax[r, c].set_title(label)
データローダーの設定
PyTorchでは、データローダー(DataLoader)を使ってデータをロードするのが一般的です。データローダーは、バッチ単位でのデータのロード、データのシャッフル、マルチスレッドによるデータの読み込みなどをサポートしてくれるので、自身でコードを書かなくても、バッチ単位で画像をシャッフルしつつ読み込むことができます。
データセットを定義していれば、データローダーの設定は非常に簡単です。
batch_size
はバッチサイズ、shuffle
はシャッフルするかどうかのフラグです。訓練データは毎回順番が変化した方がよいので、shuffle=True
と設定しています。
batch_size = 64
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
PyTorchのデータセットは、トレーニングやテストに使用するデータを扱うためのもので、例えば画像やテキストなどのデータを1つずつ処理します。一方、データローダーは、データセットからバッチ単位でデータを取り出すためのものです。この2つを組み合わせることで、簡単に大量のデータを読み出すことができます
以上で、データセットを読み込む準備は完了です。
モデル
モデルを定義する
MNISTは簡単なモデルでも学習できるので、今回は全結合層3つからなるモデルとしてみます。
PyTorchでモデルを作成する場合は、nn.Module
を派生したクラスを定義します。基本的には、以下の2つの関数を定義します
__init__()
初期化を行う関数です。引数は、自身のモデルに合わせて変更します。今回は、input_size(=28)を引数として受け取る想定です。画像のサイズはinput_size*input_size
となります。また、3つのnn.Linear
を定義していますforward()
モデルの流れ(処理)を記述します。x
は入力です(model(x)
で渡される引数)
class MyModel(nn.Module):
def __init__(self, input_size):
super(MyModel, self).__init__()
self.size = input_size*input_size
self.fc1 = nn.Linear(self.size, 1024)
self.fc2 = nn.Linear(1024, 256)
self.fc3 = nn.Linear(256, 10)
def forward(self, x):
x = x.view(-1, self.size)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
モデルを図示すると以下のようになります。最後の出力の10がクラスに対応し、クラスに対応してどれか1つの値だけ大きな値となります。例えば、文字が4の場合は、4番目の値が最大になるように学習させます。
畳み込みニューラルネットワーク(CNN)の例は以下の記事にあります
モデルを生成する
モデルを定義したので、モデルを生成します。モデルの生成は以下のコードになります。
device = ...
の行は、GPUがあるかどうかをチェックして、GPUがある場合はcuda
を、無い場合は、cpu
を設定します。
to(device)
とすることで、GPUがある場合はGPUにモデルがロードされます。
Macの場合は、GPUはmps
となります
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MyModel(28).to(device)
model
以下がモデルの内容を表示したものになります
MyModel(
(fc1): Linear(in_features=784, out_features=1024, bias=True)
(fc2): Linear(in_features=1024, out_features=256, bias=True)
(fc3): Linear(in_features=256, out_features=10, bias=True)
)
学習
損失関数と最適化手法を設定
損失関数と最適化手法を設定します。
クラス分類なので、損失関数はCrossEntropyLoss
を利用します。最適化手法は、SGDやAdamなどありますが、ここではAdamを選びました。
最適化関数の引数は、最適化するパラメータのリストと最適化のためのハイパーパラメータとなります。今回は、lr
だけ設定し、他はデフォルトを利用しています。
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
最適化のパラメータを調整することで、精度や学習速度が変化します。ここは、何度か実験して決めていくのが一般的です
訓練用と評価用関数の定義
1エポック分の訓練用と評価用の関数を用意します。これを用意しておくと、学習コードが簡単になりますので、作成するように習慣づけしておくと良いと思います。
訓練用関数
学習では、model.train()
でモデルを学習できる状態に切り替えます。
for
文で、データローダーからデータを読み込み終わるまで繰り返します。処理の内容は以下の通りです
- 画像とラベルを読み出す
model(images)
で、推論結果を出力loss
を計算loss.backward()
を実行して誤差を逆伝播optimizer.step()
を実行- 各種結果集計
どのような学習でも基本の流れはほぼ同じです。
なお、model
を呼び出す前に、optimizer.zero_grad()
を呼び出すのを忘れないようにしましょう。
def do_train(model, device, loader, criterion, optimizer):
model.train()
tot_loss = 0.0
tot_score = 0.0
for images, labels in tqdm(loader, desc="train"):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
tot_loss += loss.detach().item()
tot_score += accuracy_score(labels.cpu(), outputs.argmax(dim=1).cpu())
tot_loss /= len(loader)
tot_score /= len(loader)
return tot_loss, tot_score
評価用関数
評価では、学習を行わないのでmodel.eval()
で評価モードに切り替えます。
また、with torch.no_grad()
で勾配計算を無効化しておきます。
no_grad()
としておくことで推論時に不必要なメモリ使用量を減らすことができます
学習を行わないので、optimizer
とloss.backword
は必要ありません。
def do_valid(model, device, loader, criterion):
model.eval()
tot_loss = 0.0
tot_score = 0.0
with torch.no_grad():
for images, labels in tqdm(loader, desc="valid"):
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
tot_loss += loss.detach().item()
tot_score += accuracy_score(labels.cpu(), outputs.argmax(dim=1).cpu())
tot_loss /= len(loader)
tot_score /= len(loader)
return tot_loss, tot_score
学習ループ
学習ループです。エポック数は10回としました。
処理としては、訓練→評価を繰り返して実行しているだけです。
train
とvalid
で異なるデータを与えています。valid
のデータは学習に含まれていないので、valid
に対する結果は未知のデータに対する結果と考えることができます
valid
の結果が良いモデルを保存などすると、学習には使っていないですが、間接的にvalid
がリークするので注意が必要です。
num_epochs = 10
for epoch in range(num_epochs):
print(f'[EPOCH {epoch+1}]')
train_loss, train_acc = do_train(model, device, train_loader, criterion, optimizer)
valid_loss, valid_acc = do_valid(model, device, valid_loader, criterion)
print(f"--> train loss {train_loss}, train accuracy {train_acc}, valid loss {valid_loss} valid accuracy {valid_acc}")
最終的に、以下のような結果となりました。
--> train loss 0.014326131248904268, train accuracy 0.9952192164179104, valid loss 0.10105971224019442 valid accuracy 0.9776074840764332
評価データで97.76%とかなり高い精度で数字が認識できていることが分かります。
UMAPで特徴量を見てみる
番外編として、UMAPで学習した結果を見てみます。確認するのは以下の図の赤丸で囲んだ部分です。
最後の出力は、10クラスの分類に対応するので、1つ前の層には10クラスに分類するための特徴量が入っていると考えられます。
この部分は、深層距離学習に関係する部分です。MNISTの学習とは直接関係ないので、興味がなければ読み飛ばしてください
以下のコードではmodelのfc3
をnn.Identity
に書き換えて、評価データを流し特徴量を取り出しています。nn.Identity
は入力をそのまま出力する層で、結果としてfc2の出力結果がモデルから出力されるようになります。
model.fc3 = nn.Identity()
model.eval()
features = None
classes = None
for images, labels in tqdm(valid_loader):
with torch.no_grad():
images = images.to(device)
outputs = model(images)
# print(outputs.shape)
if classes is None:
classes = labels.cpu()
else:
classes = torch.cat((classes, labels.cpu()))
if features is None:
features = outputs.cpu()
else:
features = torch.cat((features, outputs.cpu()))
UMAPをインストール
UMAPをインストールしていない場合はインストールします。
コマンドラインでpip install umap-learn
とするか、Google Colabのセルで以下のコードを実行します。
!pip install umap-learn
可視化
以下は、UMAPで256次元を2次元に圧縮し、可視化するコードです。
import umap
umap = umap.UMAP(n_components=2, random_state=42)
X_umap = umap.fit_transform(features)
plt.scatter(X_umap[:, 0], X_umap[:, 1], c=classes)
plt.show()
色は数字別に割り当てられており、同じ色が同じ数字となります。
結果を見ると、ところどころミスはありますが、結構うまく分離できていることが分かります。
このように、特徴量を可視化すると、モデルがクラスをどの程度分離できているのかを確認することができます
おまけ:モデルをCNN(resnet18)に変更
resnet18のモデルを生成
今回は、簡単なモデルを自作しましたが、TIMM(PyTorch Image Models)などを使うともっと複雑なモデルを手軽に利用することが可能です。
timmをインストールする必要があります。timmのインストールと使い方については以下の記事を参照してください
具体的な手順は以下の通り:
- 「モデルを定義する」で説明したMyModelは削除
- 「モデルを生成する」の部分を以下のように書き換え
import timm
device = "cuda" if torch.cuda.is_available() else "cpu"
model = timm.create_model('resnet18', num_classes = 10, in_chans = 1).to(device)
model
以上で、入力が1チャンネル、クラス分類数が10クラスに調整したresnet18を利用することができます。
なお、これ以外の部分は、変更する必要はありません。
学習結果
このモデルの学習結果は以下のようになります。
--> train loss 0.01648723194505562, train accuracy 0.9949360341151386, valid loss 0.031225692634458624 valid accuracy 0.9919386942675159
最初のモデルは評価データで97.76%でしたが、resnet18では99.19%と精度が向上していることが分かります。
実際、UMAPの特徴量を見ても、各クラス間が大きく離れていることが分かります。
resnet18版のGoogle Colabで動作するコードはこちらに置いています
まとめ
MNISTの手書き数字データを識別するモデルをPyTorchを使って作成し、学習させてみました。次は、これを基本モデルとして、深層距離学習にチャレンジしてみたいと思います。
MNISTのモデルをベースにした深層距離学習の記事は以下になります