OpenCV4Python09:OpenCV(numpy.ndarray)とPyTorchで画像分類

【０】はじめに

前回TensorFlowを使ったので、今回はTorch(PyTorch)でもを画像分類を実施してみる。

【１】MNIST用のモデル作成と保存

TorchにもPyTorch Hubという学習済みモデルのリポジトリはあるが、MNIST用の学習済みモデルはおいていない。そこで、公式サンプルを使ってMNIST用の画像分類モデルを作成し、保存することにする。

pytorch/examples

このサンプルをColaboratory(juptyer)上で動かしてモデル学習・保存して手元にダウンロードする。

まずはソースコードをColaboratory上にコピペする。

■Colaboratory上での実行のさせ方
このサンプルはターミナル上で「pythonコマンド」＋「引数」から動かすようになっている。
Colaboratory上で実行させる（セルの実行ボタン押下して動かす）には、「main()内」の引数の処理を少し書き換える必要がある。

【例1】：Colaboratory上で実行させるために書き換え

... ...
def main()
   # Training settings
   parser = argparse.ArgumentParser(description='PyTorch MNIST Example')

   ... ...

   # コメントアウトする部分
   #args = parser.parse_args()
   
   # 追記する：juptyer用コマンド引数設定
   args = parser.parse_args(['--save-model'])


... ...

▲要はターミナルから実行したときの引数パース処理で得られる文字列を直接埋め込んでいる。

また、このまま実行してもよいがGPUを使った方が早いので使用することにする。

実行すると「mnist_cnn.pt」というモデルが出来上がるので手元にダウンロードする。

■実行後

ここまでで、MNIST用の画像分類モデルの作成と保存まで完了した。次にこの作成したモデルをロードしてMNISTの画像分類をしてみる。

【２】モデルロード

保存したモデルをロードして使用するには「torch.load」と「torch.nn.Module.load_state_dict」を使う。

torch.load — PyTorch master documentation

Module — PyTorch master documentation

PyTorchにおけるモデルのセーブとロードについての詳細は以下。

Saving and Loading Models — PyTorch Tutorials 1.7.1 documentation

簡単に言うと
・ロードするモデルのニューラルネットワークはあらためてコード上に定義する必要がある

・「torch.load」と「torch.nn.Module.load_state_dict」で保存したモデルをロードして、コード上に定義したニューラルネットワークにパラメータ値を当てはめる
ということをする。

【例２】：モデルをロードして定義したニューラルネットワークにパラメータを当てはめる

import torch
import torch.nn as nn
import torch.nn.functional as F
... ...


# ニューラルネットワークの定義(保存したモデルで使ったもの)
class Net(nn.Module):
   def __init__(self):
       super(Net, self).__init__()
       self.conv1 = nn.Conv2d(1, 32, 3, 1)
       self.conv2 = nn.Conv2d(32, 64, 3, 1)
       self.dropout1 = nn.Dropout(0.25)
       self.dropout2 = nn.Dropout(0.5)
       self.fc1 = nn.Linear(9216, 128)
       self.fc2 = nn.Linear(128, 10)

   def forward(self, x):
       x = self.conv1(x)
       x = F.relu(x)
       x = self.conv2(x)
       x = F.relu(x)
       x = F.max_pool2d(x, 2)
       x = self.dropout1(x)
       x = torch.flatten(x, 1)
       x = self.fc1(x)
       x = F.relu(x)
       x = self.dropout2(x)
       x = self.fc2(x)
       output = F.log_softmax(x, dim=1)
       return output


# 上記の「Netオブジェクト」を生成
mymodel = Net()

# モデルをロードしてNetオブジェクトにパラメータ値をあてはめる
mymodel.load_state_dict(torch.load('mnist_cnn.pt'))
## cpu利用時はデバイスを明示的に指定してロードする
#mymodel.load_state_dict(torch.load('mnist_cnn.pt',map_location=torch.device('cpu')))


mymodel.eval() # 評価モードに切替

実行結果例：※colaboratory上で実施

<All keys matched successfully>
Net(
 (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
 (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
 (dropout1): Dropout(p=0.25, inplace=False)
 (dropout2): Dropout(p=0.5, inplace=False)
 (fc1): Linear(in_features=9216, out_features=128, bias=True)
 (fc2): Linear(in_features=128, out_features=10, bias=True)
)

モデルをロードするときに気をつけるのは、「GPU環境か、CPU環境か」というところ。「モデルを保存した時の環境がGPUかCPUか、それに対してモデルを読み込むときの環境がGPUかCPUか」で書き方が変わってくる。

※不一致の場合は実行時にエラーメッセージとともに指定方法が表示される、等があるので確認しよう。

これでモデルのロードまでできたので、あとはOpenCV読み込んだ画像をこのモデルに渡せばよい。

【３】OpenCVで読み込んだ画像を渡して予測させる

今回も以下画像をOpenCVで読み込ませて分類させてみる。

【例３】：OpenCVで画像を読み込んで前処理実施

import cv2 as cv
import numpy as np

... ...

# グレースケール画像として読み込み
img = cv.imread('2.png',cv.IMREAD_GRAYSCALE)

# 白黒反転
img = cv.bitwise_not(img)

# スケール変換 0～1.0
img = img / 255.0

# データ形状変換
predict_x = img.reshape([1,1,28,28])

今回の入力データは以下のように前処理を実施する。
・画素値　→　背景：黒、文字部分：白
・データ形状：(1, 1, 28, 28)  ※(n枚, channel, H, W)の形状
つまり、モデルの学習につかったデータと同じになるように入力データも形状をあわせる。（後ほど学習に使ったデータの形状は確認する→【補足参照】）

後はこの画像データをモデルに渡して画像分類させるだけ。

【例4】：読み込んだ画像を分類する

# 評価(予測するだけなので勾配計算しないようにtorch.no_grad()設定をする)
with torch.no_grad():
  pred = mymodel(torch.tensor(predict_x, dtype=torch.float32))


#結果出力
print(pred)
print(pred.argmax().item()) 

出力結果例：tensor[2]が一番大きい値→「２」と判定

tensor([[-7.0256, -4.8041, -0.0324, -4.6330, -7.5134, -8.3922, -7.9629, -5.8985,
        -4.7616, -7.3971]])
2

【補足１】入力データ構造

　入力データはモデルの学習時に使用していたデータと同じ形状である必要がある。いきなりOpenCVで読み込んだ画像データの前処理内容を書いたが、実際はモデル学習に使用したデータの確認などをしてデータ形状を把握する必要がある。
【例５】：モデル学習時に使用したデータを確認する（colaboratory上で実施）

from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt

transform=transforms.Compose([  transforms.ToTensor() ])
test_data = datasets.MNIST('../data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False)
tmp = test_loader.__iter__()
data, labels = tmp.next() 

# size
print(data.size()) # torch.Size([1, 1, 28, 28])

# imshow
img = data.numpy().reshape((28, 28))
plt.imshow(img, cmap='gray')

実行結果：Colaboratory上で実施

【補足２】PyTorchとTensorflowの画像データの並びの違い

PyTorchとTensorflowは画像データ形状の並びが異なる。
　・PyTorch(n, channel, H, W) → channelが先
　・Tensorflow(n, H, W, channel) → channelが後
という違いがあるのでデータの形状変換時は注意しよう。

最初にPyTorch Hubというワードもでてきたので、PyTorch Hubにある学習済みモデルをつかった場合も簡単にまとめておく。が、長くなったので次回に続く。