初級ディープラーニング
記事内に商品プロモーションを含む場合があります

手書き文字(MNIST)認識をPyTorchで記述する【初級 深層学習講座】

tadanori

いまさらですが、MNISTの手書き文字データセットを学習するコードをPyTorchを使ってスクラッチから作成してみました。理由は、深層距離学習のテストコードを書くベースコードが欲しかったからです。

この記事では、MNISTの手書き文字データセットを学習コードを、久しぶりにスクラッチから作成しましたので(ベースモデルなしで作成)、解説したいと思います。

はじめに

MNISTの手書き数字のデータセットは、「0」〜「9」の数字を手書きしたデータセットです。

ディープラーニングの学習用によく利用されているため、手軽にダウンロードしてデータセットとして利用する環境が整っています。

今回は、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.notebooktqdmに変えてください。コードは以下になります。

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)
MNISTデータセットの中身(一部)

データローダーの設定

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のデータセットとデータローダーの関係

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番目の値が最大になるように学習させます。

モデルイメージ

モデルを生成する

モデルを定義したので、モデルを生成します。モデルの生成は以下のコードになります。

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()としておくことで推論時に不必要なメモリ使用量を減らすことができます

学習を行わないので、optimizerloss.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回としました。

処理としては、訓練→評価を繰り返して実行しているだけです。

trainvalidで異なるデータを与えています。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}")
Google Colabでの実行結果

最終的に、以下のような結果となりました。

--> 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のfc3nn.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()

色は数字別に割り当てられており、同じ色が同じ数字となります。

結果を見ると、ところどころミスはありますが、結構うまく分離できていることが分かります。

UMAPの結果

このように、特徴量を可視化すると、モデルがクラスをどの程度分離できているのかを確認することができます

おまけ:モデルをCNN(resnet18)に変更

resnet18のモデルを生成

今回は、簡単なモデルを自作しましたが、TIMM(PyTorch Image Models)などを使うともっと複雑なモデルを手軽に利用することが可能です。

timmをインストールする必要があります。timmのインストールと使い方については以下の記事を参照してください

PyTorch-TIMMでモデル作成、モデル一覧を取得する方法(create_modelチートシート)
PyTorch-TIMMでモデル作成、モデル一覧を取得する方法(create_modelチートシート)

具体的な手順は以下の通り:

  • 「モデルを定義する」で説明した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を利用することができます。

ResNet18は、畳み込みニューラルネットワーク(CNN)モデルの1つです。”ResNet”は、”Residual Network” の略称であり、その名前はこのモデルの特徴である「残差学習」(residual learning)に由来しています

なお、これ以外の部分は、変更する必要はありません。

学習結果

このモデルの学習結果は以下のようになります。

--> train loss 0.01648723194505562, train accuracy 0.9949360341151386, valid loss 0.031225692634458624 valid accuracy 0.9919386942675159

最初のモデルは評価データで97.76%でしたが、resnet18では99.19%と精度が向上していることが分かります。

実際、UMAPの特徴量を見ても、各クラス間が大きく離れていることが分かります。

UMAPの結果(resnet18)

resnet18版のGoogle Colabで動作するコードはこちらに置いています

まとめ

MNISTの手書き数字データを識別するモデルをPyTorchを使って作成し、学習させてみました。次は、これを基本モデルとして、深層距離学習にチャレンジしてみたいと思います。

MNISTのモデルをベースにした深層距離学習の記事は以下になります

深層距離学習+FaissをMNISTデータセットで実践【Pytorch Metric Learning】
深層距離学習+FaissをMNISTデータセットで実践【Pytorch Metric Learning】

おすすめ書籍

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

記事URLをコピーしました