画像分類モデルを自作してみた & Web上から触れるようにしてみた
この記事はブログ記事 (ふりかえり) 兼 課題発表用記事として作成しています。
課題制作したシステムの大まかな説明と、AIモデル自作に関するそこそこ詳細な説明が含まれます。
また、多数の素人理解が含まれます (特にニューラルネットワーク辺り) 。
前々からPyTorch, torchvisionを利用した、画像を扱う何かしらを作ってみたい欲求はあったものの、何かと理由をつけて逃げていましたが、学校の課題でAIモデルを扱った何かを作る必要が出てきて「せっかくだしな」ということで。
VRSNS「VRChat」の3Dアバター35体の顔写真をもとに画像分類AIモデルを作ってみました。
若干過学習気味かつデータセットの問題で汎用性は低いですが、サンプルアバターやBOOTH ショップ画像ぐらいならおおよそ正しく推論できるっぽいので及第点とします。
厳密には1からの自作ではなく、 ResNet50.IMAGENET1K_V2
をベースとした転移学習です。が、ここでは便宜上「自作」として扱います。そのほうがわかりやすいので![1]
この記事を書いている裏でも学習が回っています。
1. おおざっぱな概要
PyTorch 1.13.1 + Torchvision 0.14.1 + ResNet50.IMAGENET1K_V2 を利用して、3Dモデル合計1750枚の写真を使って画像分類モデルの転移学習を行いました。
フロントエンド/バックエンドは Flask を利用しており、Web上から画像をアップロードしたり、torchを呼び出してAIによる推論を行えるようになっています。
プロジェクトは uv を用いてPythonバージョン・パッケージ管理を行っています。
また、ソースコードの大部分は ml-flaskbook/flaskbook をベースとしています。
AIモデルのトレーニングに使ったデータセットは全体合計で1750枚、うちモデル1体 (1クラス) につき50枚使用しており、35クラスあるため 35 * 50
より1750枚となります。
画像分類システムの概要
ユーザー作成またはログイン後、画面右上「画像新規登録」ボタンから画像をアップロードし、アップロードした画像の「検知」ボタンを押すことでtorchが呼び出され、AIモデルによる画像の分類処理が行われます。
それだけです。画像の削除とかもできますが、実ファイルが消えるわけではない、いわゆる論理削除です。
また、複数体写ってる画像でもある程度の精度は出ます。流石に違うキャラクターが混ざってるとおかしくなりやすいですが。
ただしこのタグ表示の部分は改良の余地ありかなと思います。そもそも画像分類モデルで物体検出っぽいことをしようとしてるのが間違いなんですけど。
また、画像の削除や検知はログインしていないとできません。
2. AIモデルについて
本題。
学習環境
- CPU (学校用PC): Intel Core i5 1340P@1.9Ghz
- GPU (メインPC): NVIDIA GeForce RTX 2060 6GB (Mobile)
- GPU (学校PC): NVIDIA GeForce RTX 3070 8GB
当初はCPUで回していましたが、GPUで回せる方法を得てからは2060、開発終盤は基本3070を回していました。
Google Colab GPUなどのクラウド環境での学習も試してみたいところではあります。無料の範囲でもTesla K80とか回せるらしい。金払うとV100とかA100も回せて夢がある。
データセット
今回使ったデータセットの大元の3Dモデルは、ECサイト「BOOTH」の「3Dキャラクター」カテゴリ「人気順」ソートで1ページにある上位15体かつ適当にピックアップした20体の計35体の3Dモデルを使っています。
これらの商品ページの画像や、実際の3Dモデルの画像を撮影して1体あたり50枚の画像を用いてトレーニングとテストを行っています。
これらのモデルは本来「VRChat」で使われることを想定した3Dモデルです。
普通の3Dモデルと違い、細部まで作りこまれていることや、VRChatやその他3D制作ソフトなどに持ち込みやすいデータ構造をしているため、データセットの収集が非常に楽でした。
ただし、各3Dモデルの利用規約には注意が必要だと思います。今回選定したモデルは基本的にAI学習禁止などは見受けられませんでしたが、モデルによってはそういった規約が存在する可能性もあります。
データセットの管理自体も不適切でした。
当初は多くて100枚ぐらいを想定していたため、Git リポジトリ内にそのまま置いてGitHubに上げていましたが、データセットを増やす段階になると合計で1000枚を優に越してしまい、かなり不適切なリポジトリの状態にしてしまいました。
最初の設計段階からHuggingFace等を活用すべきでした。使って損することはないと思うので。
3. モデル自作の歴史
現在のモデルに落ち着くまでの開発の大まかな道のりです。
転移学習に用いたコードはQiitaやらGemini生成物やらの寄せ集めでとても公開できるものではありません。が、今見るとあまり綺麗なものではないように思います。
いつか書き直したい。
初期段階
開発初期時点ではまだクラスが2体しかなかった上、そもそものデータセット数が非常に少なく (大体10枚ぐらいだったような) 、データ拡張[3]の設定自体も細かく調整していないためにとても精度が低く、実用的ではありませんでした。
なんならあまり学習も回せていません。90MBのゴミ。
その後、データセット量を増やしたりデータ拡張パラメーターを調整することで大きく精度を上げられ、なんとか実用可能なところまで性能を上げることに成功しました。
GPUを扱えるように改良したため、学習回数10エポック未満からなんとか25~50エポックぐらいまでは回せるようになりました。
学習にGPUを使えるようになったのが大きいと思います。CPUだとたった1イテレーションに1秒とかザラにかかってしんどかった。2060なら約13イテレーション毎秒!純粋に13倍!3070ならなんと約30イテレーション!
やはり数の暴力こそが正義。CUDAコアで殴れ!
Flaskと連携
学習し終えたなら次はWebから扱えるよう、Flaskに組み込まなければなりません。課題として成立しないので。[4]
ただ大元のソースで指定するモデルファイル名を変えたりしただけでは動かず。
当然です。もともとのソースで使っていたモデルである ResNet50.IMAGENET1K_V2
は1000層もの出力層を持っていますが、このモデルはたった2層しか持っていません。
そのため出力層などからの違いから、そのままではPyTorchにロードすることができません。
元のコード (とモデル) では特に何も考えず扱えるよう最初からいい感じになるようになっているようですが、この特製モデルはそんな親切なことはしていないので、読み込むときにモデルの詳細情報をPyTorchに教えてあげなければいけません。
ということでさっくり教えます。これで本当に適切なのか知りませんが動いてるのでよし。StackOverFlowいつもありがとう。
特に model.load_state_dict
が必須。torch.load()
でいけるだろうと思っていたがダメだった・・・理解と学習が足りない・・・
モデル自体は正しく動くようになり、推論もうまく動くようになりました。
ではこれで終わりなのかというとそうではなく。ここまでうまく動いたならもっと色々やりたくなってしまった。
無限改良編
このままでは224x224の画像でないと正しく推論が走りません。
大本のモデルの ResNet50.IMAGENET1K_V2
は224x224の画像でトレーニングされているため。転移学習を行う際も同じく224x224の画像データを使わなければいけないためです。
そして大体のAIはトレーニングに使った解像度と同じ解像度じゃないとうまく動きません。[5]
ということで、「スライディングウィンドウ法」という方法を用いて、高解像度の画像データでも正常に推論ができるようにしました。なお厳密なスライディングウィンドウ法の定義とはズレています
やってることは下の折りたたみ開いたらでてきます。
端的に言うと、「大きな配列を適当に区切って、そこに対して推論をかけたあと配列をずらしてまた推論をかける手法」のことです。
が、今回の実装内容は厳密にはこれではなく、「大きな画像を一定サイズで区切って1枚1枚に推論を走らせる方法」を取っています。
詳細はこちらのQiitaがわかりやすいです → スライディングウィンドウ法とは?
やり方は近いのでセーフということで。
このあたりのコーディングは半ばGemini依存みたいなところがあります。あまりよろしくないけど、処理の手順の流れ自体は容易に把握できるので「コピペプログラマー」とは違う・・・とは思いたいなあ。
この方法の利点はわざわざリサイズ処理を挟んだりしなくとも、大きな解像度の画像をそのまま推論できるところにあります。
画像によっては潰れてしまったりすることがあるので、見た目を維持したまま推論できるのはいいことだと思います。
そのかわり、非常に大きな欠点として、推論したい対象 (この場合では顔) が切り出された224x224の範囲に収まらないと正常に推論できません。
どうなるのかというと、こうなります。
こればかりはどうしようもありません。克服するにはモデル自体をもっと高解像度の画像でトレーニングできるようにしなければいけませんが、ResNet50では無理です。
また、ただでさえ224x224の画像1枚に対してでもそこそこ時間のかかる処理が大量に行われるので待ち時間はかなり長いです。
具体的に言うと、i5-1340Pで推論をしているとき、電源管理を最適なパフォーマンスに設定していても1920x1080の画像に対して20秒ぐらいかかります。[6]
GPUを回せば爆速で終わるので純粋にCPUの性能が低すぎる気がします。
GPUを使っているなら model.half()
でfp16化するなどの方法で高速化できます[7]が、CPUを使ってる場合はむしろ逆に遅くなってしまう[8]ため、なかなか手も出せず。難しい。
Intel GPU向けPyTorchを使えばいいのかもしれませんが、NVIDIA GPU向けTorchと共存できないと思うので結構めんどくさそう。
現状では打つ手なしということで、GPUで回すことを想定した設計としています。課題としていいのかこれ
推論方法の最適化ができればもっと短縮できる気はしますが。
高解像度の推論ができるようになったら、次は「画像のどこがキャラクターなのか」を示すバウンディングボックスが欲しくなりました。
ということで、OpenCVを振り回して画像分類モデルなのに物体検出モデルのような挙動を実装しました。今思えばこの時点でYOLOなど物体検出に特化したモデルで再度転移学習を行った方がよかったのではないかと思います。
区切って推論した画像の中で最も「自信」が高い順に並べ、上位10個に箱を描画することでそれらしさを出しています。
ただしこのままだと線が太くて画像サイズによっては大幅に視認性が悪くなってしまいます。
なので1位のものだけ不透明で、それ以下は半透明で描画するようにしました。多少はマシになったと思います。
頑張れば画像サイズに応じて動的に線の太さを変える処理を実装できるような気もします。
ここまでできたらあとはモデルの改良だけです。この時点で締切1日前です。時間がなさすぎる。
まだまだやってみたい最適化手法などはあったので、もっとスケジュールに余裕を持ってやるべきでした。
データセットを増やす
まずは識別できるキャラクター量を増やすために、クラス (とデータセットそのもの) を増やしました。
Boothの商品ページから画像を切り出してきたり、実際にVRChat内でサンプルアバター[9]の写真を撮影したりして、とりあえず1クラスあたり20~40枚ほど用意しました。
ただ、この状態だとデータセットの不均衡が発生してしまい、思うように精度が出ませんでした。どうしてもデータセットが多いほうに推論が寄ってしまいます。
このときのデータセットの配分は「最大で43枚」「最小で22枚」と大きな差があいてるものでした。
データ拡張をしたとしてもちょっと悪影響が心配になるレベルです。(というか悪影響が出てます)
ということで、少ないほうを多くする方向でデータ量を調整しました。数えてませんがかなりの量の写真を撮影したりトリミングしたりした気がします。10時間ぐらいかかった。
あまりにファイルの変更を加えすぎて git add .
にめちゃくちゃ時間がかかりました。そもそもこういうことにGitを使ってはいけないのですが。
学習の最適化
データセットの拡張後、学習時間が伸びすぎてしまうことの対策として学習スクリプトの最適化を行いました。
具体的には、データ拡張パラメーターの変更、optmizerとschedulerの変更、学習時のみ半精度化、学習率の変更、過学習対策のCSV書き出し処理の追記を行いました。順を追って説明します。
データ拡張パラメーターの変更
水増しのやり方を変えたともいいます。
今まではすべての水増し方法の適用を100%行うようにしていましたが、学習時の悪影響を考えて50%または20%の確率で適用されるように変更し、色相を変える処理は変更幅を大きく落としました。
元画像または元画像に近いデータが多く含まれるようになるため、汎化性能の向上を期待しています。
optimizer, schedulerの変更 (失敗かも)
はっきりいって、これはやるべきではありませんでした。
Geminiに最適化手法を聞いたときに真っ先に出てきて、効果をろくに調べないままやってしまいました。AI妄信の欠点。
optimizerとは、モデルの最適化手法のことです。損失関数や勾配からパラメーターを更新してモデルの最適化を行うとき、指定されたoptimizerが働きます。
詳しくはこちら: Pytorchの様々な最適化手法(torch.optim.Optimizer)の更新過程や性能を比較検証してみた!
schedulerとは、学習中の学習率 (パラメーターの更新率) を変えられるものです。学習中ずっと同じ学習率では後半になってくると安定しなくなるため、途中で下げたりすることによって安定化を狙います。
詳しくはこちら: 【PyTorch】学習率スケジューラー解説
最初から使っていたもので問題なかったのに、Geminiのいうまま、optimizerを AdamW
、schedulerを CosineAnnealingLR
に変えてしまいました。
転移学習において、AdamWもCosineAnnealingLRも基本的には非推奨のようで。通常の学習 (?) であれば十分らしいのですが。
もともとoptimizerに SGD
、schedulerに StepLR
を使っていました。
かなり説明が難しいので詳細な説明は省きますが、SGDもStepLRも簡単に言えばほぼ直線的に収束に向かって動くようなもの、として認識しても問題ないと思います。
対するAdamWはウネウネ動いて徐々に収束に向かって動いていきます。CosineAnnealingLRはコサイン波の動きで学習率を変動させていきます。もしかしたらこれがまずかったかも。
結果として、100エポック回しても全く学習が進まず、時間ばかりかかって成績のいいモデルが得られませんでした。
過学習のような状態が続き、かといって早い段階では学習不足などといったかなり難しい状態になってしまっていました。
おとなしく元のoptimizer, schedulerに戻したところ改善したので、おそらくこのケースにおいては適さないものだったかと思います。
かといってそのまま捨てて二度と使わないのもどうにももったいないような気がするので、またいつかの勉強のタイミングで使ってデータを見てみたいところ。
もしかしたら200とか300エポックぐらい回したらいい感じに収束したのかもしれない。それはそれで使い勝手悪すぎるので結局使わないと思いますが。
学習率の変更
ちょっと時間軸は前後しますが。optimizer, schedulerの変更後、しばらく学習を回してパラメーターを探っていましたが、どうにもtrain, val両方ともloss[10], acc[11]ともに安定しない現象が発生しました。
いわゆる「過学習」や「学習不足」が起こってるような挙動が起こっていましたが、問題なのは過学習が起きるほど回していなかったり、学習不足が起きるぐらいに回していないことはないということです。
過学習とは、モデルがデータセットを「丸暗記」してしまい、未知のデータに対しての応用力を失っている状態です。
train_lossは下がり続けているのにval_lossは下げ止まるかまたは上がり始め、
train_accは100%に近づき、val_accは逆に下がり始める挙動を示します。
学習不足とはその逆で、まだデータセットを覚えていない状態です。
train, valともにlossが下がりきっておらず、accも低いまま上がりません。
学習を進めれば順調に精度が上がっていく可能性があります。
正常に学習できている場合は、データセットを適度に覚え、未知のデータに対しても問題なく推論ができる状態です。
train, valともにlossが程よく下がり続け、accも上がっていきます。
ここでlossの減少やaccの上昇が止まったり、振動するようになったら要注意のサインです。
調べてもどうにも情報が出てこず、やむなくGeminiに聞いてみたらところ「lr[12]を落としてみたらいい」という案を提示されたため、初期値の 0.001
から 0.0005
まで落としました。
結果としてはあまり変わらず。ここでoptimizer, schedulerをSGD, StepLRに戻したところ大きく改善した、というわけです。
厳密にlrを下げる前とのデータを取っていないため、学習時間や精度がどう変わったかは差を取れませんでしたが、少なくともlossが収束して安定するまでが多少早くなったような感じはします。
学習時の半精度化
通常のPyTorchは学習、推論ともに「fp32 (single-precision)」という単位を元に計算を行っていますが、使用するメモリ量や帯域は純粋に「fp16 (half-precision)」の2倍です。
「fp32」とは正式名称「IEEE 754 single-precision」です。有効桁数7桁の浮動小数点数を扱うことができ、非常に高精度の値を持つことができます。
対する「fp16」とは「IEEE 754 half-precision」です。その名の通り、fp32の半分ほどの精度しか出せませんが、使用するメモリ量も同じく半分ぐらいまで減ります。
Stable Diffusionが出てきた当初はモデルをfp16化することで推論を高速化する、などの「ハック」が有名でした。
ここで、学習時かつ計算部分のみfp16で受け持つことで、使用するメモリ量を減らしたり、学習時の高速化を狙う最適化手法を「Automatic Mixed Precision (混合精度)」と呼びます。
学習の一部分だけfp16を使うため、推論は変わらずfp32なのでCPUでもGPUでも実行できますが、当然推論時の高速化は期待できません。おそらく。
どちらかというとモデルの学習時間を短くして作業最適化を図ろうとするものです。
結論から言うと、このケースではほとんど意味がないどころか逆に学習時間が長くなりました。メモリ量もそこまで変わらず。
Tensorコアを優先的に回すようになるらしいので、Tensorコアが少ないRTX2060だと逆効果だったかもしれません。3070なら目に見える効果が期待できるところですが、もう時間がありません。
過学習を防ぐためのCSV生成
厳密には「防ぐ」ではなく「知る」ためですが。
学習時にtrain_loss, train_acc, val_loss, val_accと現在のエポック数をそれぞれ取得し、dictに押し込めて学習完了時にまとめてCSVで吐き出すスクリプトをこれまたGeminiに吐かせました。
この吐き出されたCSVを元にExcelやmatplotlibでグラフを生み出せば、エポックごとの学習状態の推測ができるというわけです。
最終的なモデル
この他にも一部細かな最適化などを行い、最終的に 35e, train_loss 0.6535, train_acc 0.8450, val_loss 0.1199, val_acc 0.9714
と割と高性能ぽさそうなモデルの学習に成功しました。
とはいえ、やはり推論方法の問題やデータセットの数の問題もあり、若干サンプルアバターにしか反応しておらず、大きく改変されたアバターの写真には間違った結果を返すことが多いようです。
こればかりはデータセットを増やすしかありません。今後も時間があったらどんどんデータセットを増やしていこうと思います。
一応 validation 用に用意した画像ならだいたい正解なんですが…なんか思ってたのと違う。
データも時間も足りない挑戦だったので、今度は余裕を持って大量のデータを集めてやってみたい所です。
4. データセットに使用した3Dモデルたち
↓ 折りたたみの中にあります! ↓
順番は名前順。
ぷらすわん - 【オリジナル3Dモデル】 アッシュ (ASH)◆セフィラ共通素体
あまとうさぎ - シフォン -Chiffon-【オリジナル3Dモデル】
あまとうさぎ - ショコラ -Chocolat-【オリジナル3Dモデル】
sep-neko-ya - 【オリジナル3Dモデル】シアン - Cian #Cian3D
VERMILION Studio - 【デルタフレア】オリジナル3Dモデル
STUDIO JINGO - オリジナル3Dモデル「イヨ」ver1.08
あまとうさぎ - カリン -Karin-【オリジナル3Dモデル】
ポンデロニウム研究所 - オリジナル3Dモデル「桔梗」
もち山金魚 - キプフェル Kipfel / オリジナル3Dモデル
ひゅうがなつみかん - 【オリジナル3Dモデル】リーファ
あまとうさぎ - ライム -Lime-【オリジナル3Dモデル】
ぷらすわん - 真冬 Mafuyu / オリジナル3Dモデル
もち山金魚 - まめひなた Mamehinata / オリジナル3Dモデル
STUDIO JINGO - オリジナル3Dモデル「マヌカ」ver1.02
IKUSIA - オリジナル3Dモデル「真央」-mao-
BeroarN - Marycia マリシア - オリジナル3Dモデル #Marycia3D
キュビクローゼット - オリジナル3Dモデル「舞夜」Ver.1.02.2
あまとうさぎ - ミルク Re -Milk Re-【オリジナル3Dモデル】
DOLOS art - オリジナル3Dモデル『ミルティナ』
あまとうさぎ - ミント -Mint-【オリジナル3Dモデル】
ポンデロニウム研究所 - オリジナル3Dモデル「ミーシェ」
IKUSIA - オリジナル3Dモデル「瑞希」メニューギミック搭載
キュビクローゼット - オリジナル3Dモデル「萌」Ver.1.02.1
rio3d - 凪 -Nagi-【オリジナル3Dモデル】
ぷらすわん - 【オリジナル3Dモデル】Nayu - ナユ
ぷらすわん - 【オリジナル3Dモデル】Platinum - プラチナ
STUDIO JINGO - オリジナル3Dモデル「竜胆」ver1.08
IKUSIA - オリジナル3Dモデル「りりか」-ririka-
IKUSIA - サメっ子オリジナル3Dモデル「rurune」-ルルネ-
あまとうさぎ - ラスク -Rusk-【Mobile/PC対応 オリジナル3Dモデル】
STUDIO JINGO - オリジナル3Dモデル「セレスティア」ver1.02
ポンデロニウム研究所 - オリジナル3Dモデル「しなの」
Chocolate rice - 【オリジナル3Dモデル】 Sio / しお / ver.2.01
もち山金魚 - うささき Usasaki / オリジナル3Dモデル
ひゅうがなつみかん - 【オリジナル3Dモデル】ウルフェリア
- ゆくゆくは完全な自作もしてみたいがめんどくさそうだなぁ・・・ ↩
- trainに40枚、valに10枚。 ↩
- いわゆる水増し。1枚の画像に対していろんなフィルタを通したりすることで5枚に増やしたりする。 ↩
- 実はAIモデルを訓練する課題ではない。Flaskを使うついでにAIを触る課題なのである・・・ ↩
- この辺は初期のStable Diffusionとかでも近かったような気がする。512x512じゃないとうまく生成できないみたいな。 ↩
- 1920x1080といえば一般的なVRChat 写真サイズ ↩
- これStable Diffusionで見た。 ↩
- 大体のCPUはFP16による演算をサポートしていないため。 ↩
- モデル購入前に実際にゲーム内で「試着」できるもの ↩
- 損失。低ければ低いほど不正解率が低く、よい。 ↩
- 精度。高ければ高いほど正解率が高く、よい。 ↩
- Learning Rate。学習率。 ↩
- 大体7桁ぐらいの有効桁数。 ↩