MoZuMaにおいてダウンロード済みのモデルファイルを使用する
はじめに
先日、「顔画像の品質を評価できるMagFaceをローカルGPU環境で動かしてみた」という記事を書きました。
MoZuMaというライブラリを使って、MTCNNで顔検出、MagFaceで顔特徴量抽出を行うことができましたが、モデルのロードがどうにも遅い…。
軽く調べてみると、ダウンロードしたモデルファイルをキャッシュせず、毎回モデルファイルをダウンロードしていることが分かりました。今回はこれを改善してみます。
モデルファイルを事前にダウンロードする
修正前のコードは以下の通りです。torch_mtcnn関数、torch_magface関数でそれぞれMTCNN、MagFaceのモデルファイルをダウンロードし、ロードしています。
import torch
from mozuma.models.magface.pretrained import torch_magface
from mozuma.models.mtcnn.pretrained import torch_mtcnn
device = torch.device("cuda")
mtcnn = torch_mtcnn(device=device)
magface = torch_magface(device=device)
デバッガを使って調べてみると、それぞれ以下のURLからモデルファイルをダウンロードしていました。
- MTCNN: https://github.com/mozuma/mozuma/releases/download/state.pytorch.facenet-mtcnn/facenet.state.gzip
- MagFace: https://github.com/mozuma/mozuma/releases/download/state.pytorch.magface/magface.state.gzip
URLが判明したので、モデフファイルを事前にダウンロードしておきます。
拡張子をgzipからgzに変更しているのは、gzipのままだとgzip -dに失敗するためです。
wget -O facenet.state.gz https://github.com/mozuma/mozuma/releases/download/state.pytorch.facenet-mtcnn/facenet.state.gzip
wget -O magface.state.gz https://github.com/mozuma/mozuma/releases/download/state.pytorch.magface/magface.state.gzip
また、ロードの度に展開するのはCPUコストが勿体ないので、事前に展開しておくことにしました。
gzip -d facenet.state.gz
gzip -d magface.state.gz
参考に、展開前後のモデルファイルサイズを以下に示します。事前に展開することによるストレージコストは微々たるものです。
| モデル | 展開前 | 展開後 |
|---|---|---|
| MTCNN | 1.8 MB | 1.9 MB |
| MagFace | 244.2 MB | 261.1 MB |
さすがに200MBを越えるモデルファイルを毎回ダウンロードするのは無駄ですね。モデルのロードが遅いはずです。
それにしてもMTCNNのモデルってこんなに小さいのですね。驚きです。
ダウンロード済みのモデルファイルをロードする
MoZuMaのソースコードを読んでみるとTorchModelクラスのset_stateメソッドが使えることが分かりました。
set_stateメソッドを使い、以下のコードでダウンロード済みのモデルファイルを読み込むことができました。
from pathlib import Path
import torch
from mozuma.models.magface.modules import TorchMagFaceModule
from mozuma.models.mtcnn.modules import TorchMTCNNModule
device = torch.device("cuda")
mtcnn = TorchMTCNNModule(device=device)
mtcnn.set_state(Path("facenet.state").read_bytes())
magface = TorchMagFaceModule(device=device)
magface.set_state(Path("magface.state").read_bytes())
おわりに
モデルファイルを事前にダウンロードし、ダウンロード済みのモデルファイルをロードすることでロード時間を改善することができました。
他の機械学習ライブラリと同じように、MoZuMaにもキャッシュ機構を内蔵してもらえると嬉しいですね。
本記事が何らかの参考になれば幸いです。
Discussion