前回に引き続きPyTorch 1.9のリリースノートを眺めていたら、PyTorchのモデルとコードをパッケージングするための新機能が入っているようだったので調べてみました。
New APIs to optimize performance and packaging for model inference deployment
前半の optimize performance の部分については、Inference Mode の紹介を前回したのでそちらを参照してください。
今回は packaging for model inference deployment の部分を調べてみたところ、torch.package というモジュールがβとして追加されたようです。
torch.package is a new way to package PyTorch models in a self-contained, stable format. A package will include both the model’s data (e.g. parameters, buffers) and its code (model architecture). Packaging a model with its full set of Python dependencies, combined with a description of a conda environment with pinned versions, can be used to easily reproduce training. Representing a model in a self-contained artifact will also allow it to be published and transferred throughout a production ML pipeline while retaining the flexibility of a pure-Python representation.
お試し
公式のチュートリアルでは、DCGANのモデルをPyTorch Hub(torch.hub)から取得してモデルをエクスポート/インポートしていますが、生成されるパッケージの内部構造等についてわかりにくかったので、手元でシンプルなモデルを作ってパッケージングします。
前準備
パッケージングのテストなので、モデルを作る部分は簡単に済ませます。ここではAlexNetをPyTorch Hubやtorchvision経由で取得せずに自前で実装し、CIFAR-10データセットを学習したモデルを準備しておきます。
- データセット: CIFAR-10
- モデルアーキテクチャ: AlexNet
学習処理部分のコードを全て載せると長くなるので、GitHubにGoogle Colab用のノートブックを上げておきます。ちなみにノートブック内では前回紹介した torch.inference_mode を使ってバリデーションとテストしてますが特に不具合などは無さそうでした。v1.9以降は torch.no_grad の代わりに使っても良さそうです。細かい学習処理部分のパラメータなどはupしたノートブックの中を参照してください。
AlexNetはCIFAR-10用に畳み込み層のカーネルサイズとか微修正したものになります。
# alexnet.py
import torch
from torch import nn
class AlexNet(nn.Module):
features: nn.Sequential
classifier: nn.Sequential
def __init__(self, num_classes: int=10):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(64, 192, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2),
)
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 2 * 2, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.features(x)
x = x.view(x.size(0), 256 * 2 * 2)
x = self.classifier(x)
return x
モデルを定義したファイルは後のパッケージング処理で必要となるのでファイル名はわかりやすいように付けておくと良いです。
モデルのパッケージング
やっと本題ですが、パッケージング作業自体は簡単です。以下のように作ったモデルとコードを torch.package を使ってパッケージングします。
import torch
from torch import package
from alexnet import AlexNet ## 前述のAlexNet (alexnet.py)
model_path = "./cifar10_net.pth" ## 前述のAlexNetでCIFAR-10を学習したモデルファイルへのパス
pt_path = "./alexnet_cifar10.pt" ## 任意の出力パッケージファイルのパス
package_name = "alexnet_cifar10" ## 任意のパッケージ名
resource_name = "model.pkl" ## パッケージ内に含める任意のモデルファイル名 (pickleファイル)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = AlexNet()
model.load_state_dict(torch.load(model_path, map_location=device))
print(model)
## torch.package.PackageExporter でパッケージング処理
with package.PackageExporter(pt_path, verbose=False) as exporter:
exporter.intern("alexnet") ## 前述のalexnetモジュール(alexnet.py)をパッケージに含める
exporter.save_pickle(package_name, resource_name, model) ## パッケージファイル出力
これで alexnet_cifar10.pt という名前のパッケージファイルが出力されます。PyTorch公式には拡張子は.ptを使っているようで、ファイルの実体はzipファイルとなっています。
$ file alexnet_cifar10.pt alexnet_cifar10.pt: Zip archive data
次はこのzipファイルの内部構造を見てみます。
$ unzip alexnet_cifar10.pt
$ tree -a alexnet_cifar10
alexnet_cifar10
├── .data
│ ├── 94579536640256.storage
│ ├── 94579536682800.storage
│ ├── 94579537130400.storage
│ ├── 94579537131952.storage
│ ├── 94579537132464.storage
│ ├── 94579537134640.storage
│ ├── 94579537135152.storage
│ ├── 94579537136816.storage
│ ├── 94579537137328.storage
│ ├── 94579537138944.storage
│ ├── 94579537141824.storage
│ ├── 94579537158848.storage
│ ├── 94579537159360.storage
│ ├── 94579537176384.storage
│ ├── 94579537340864.storage
│ ├── 94579537341536.storage
│ ├── extern_modules
│ └── version
├── alexnet.py
└── alexnet_cifar10
└── model.pkl
今回はシンプルなモデルなので、上記のパッケージング処理のコードを見れば内部構造との対応がだいたいわかるかと思います。指定したパッケージ名(alexnet_cifar10)のディレクトリが作られ、その中に model.pkl ファイルが出力されています。また、AlexNetを定義したファイル(alexnet.py)も指定通りに含まれています。補足として、.data ディレクトリ内にある .storage ファイル群はシリアライズされたTensorデータの実体となるバイナリファイル、extern_modules は依存モジュールが書かれたテキストファイルです。依存モジュールは良い感じに自動検知してくれます(内部的にはpickletoolsを使用して依存関係チェックしている)。
$ cat alexnet_cifar10/.data/extern_modules torch torch.nn collections builtins torch.nn.modules.container torch.nn.modules.conv torch._utils torch.nn.modules.activation torch.nn.modules.pooling torch.nn.modules.dropout torch.nn.modules.linea
torch.package は内部的には pickle モジュールに依存しており、モデルのファイル形式もpickleフォーマットになります。
モデルの読み込み
読み込みも簡単です。torch.package.PackageImporter を使います。
pt_path = "./alexnet_cifar10.pt"
package_name = "alexnet_cifar10"
resource_name = "model.pkl"
## torch.package.PackageImporter でパッケージ読み込み
importer = package.PackageImporter(pt_path)
loaded_model = importer.load_pickle(package_name, resource_name)
print(loaded_model)
## 出力結果
AlexNet(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(1): ReLU(inplace=True)
(2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(3): Conv2d(64, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): ReLU(inplace=True)
(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(7): ReLU(inplace=True)
(8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(9): ReLU(inplace=True)
(10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(classifier): Sequential(
(0): Dropout(p=0.5, inplace=False)
(1): Linear(in_features=1024, out_features=4096, bias=True)
(2): ReLU(inplace=True)
(3): Dropout(p=0.5, inplace=False)
(4): Linear(in_features=4096, out_features=4096, bias=True)
(5): ReLU(inplace=True)
(6): Linear(in_features=4096, out_features=10, bias=True)
)
)
正しく読み込みができているようです。
機能説明補足
torch.package モジュールについて機能面を整理してみます。ここでは簡単なモデルをパッケージングしたのでモデルファイル以外の添付物はalexnet.pyのみでしたが、テキストファイルなど任意のファイルを含めることができます。torch.package.PackageExporter には save_pickle 以外にも save_text (テキストファイルを添付)や save_binary (Pythonオブジェクトをシリアライズしたバイナリファイルを添付)メソッドが提供されています。torch.package.PackageImporterによる読み込みも同様です。
exporter.save_pickle("my_resources", "tensor.pkl", torch.randn(4))
exporter.save_text("config_stuff", "words.txt", "a sample string")
exporter.save_binary("raw_data", "binary", my_bytes)
また、パッケージ内に含める/含めないは intern / extern メソッドを使います。globパターンでの指定もOKです。注意点として、NumPyなどのC拡張モジュールはパッケージには含められないので明示的に extern で外しておく必要があります。また、モデル作成の再現ができるように学習部分の実装コードも含めておくと良いでしょう。
## torchvision.* モジュールを全て含める
exporter.intern("torchvision.**")
## numpy モジュールを含めない
exporter.extern("numpy")
その他細かい機能はいろいろあるのですが、しばらく使ってみて知っておいた方が良い機能を見つけたらまた紹介します。
不明、懸念点など
上記サンプルコードを見てもわかるように、torch.package.PackageImporter でパッケージを読み込む場合、モジュール名やモデルファイル名を認識していないと読み込めないことになります。パッケージファイル単体で管理することは想定されていないということでしょうか。
ちなみにパッケージ内部構造はzipファイルを展開しなくても、以下のようにAPI経由で一応確認は出来ます。treeコマンドのような出力になるようです。ただ、内部構造がわかったところで、model.pklファイルがモデルファイルかどうかはわからないのですが。パッケージ内はモジュール毎にディレクトリが分かれているので、一つのモジュールにモデルファイルが一つだけという規約を作ればモデルファイル名は固定(例えばmodel.pkl)してもいいのかなとは思います。
importer = package.PackageImporter(pt_path)
print(importer.file_structure())
## 出力結果
─── ./alexnet_cifar10.pt
├── .data
│ ├── 93997045004704.storage
│ ├── 93997045047360.storage
│ ├── 93997045494960.storage
│ ├── 93997045496512.storage
│ ├── 93997045497024.storage
│ ├── 93997045499200.storage
│ ├── 93997045499712.storage
│ ├── 93997045501376.storage
│ ├── 93997045501888.storage
│ ├── 93997045503552.storage
│ ├── 93997045506432.storage
│ ├── 93997045523456.storage
│ ├── 93997045523968.storage
│ ├── 93997045540880.storage
│ ├── 93997045540992.storage
│ ├── 93997045706048.storage
│ ├── extern_modules
│ └── version
├── alexnet_cifar10
│ └── model.pkl
└── alexnet.py
懸念点としては、この機能は1.9から追加されたということで、1.8以前では当然動かないということになります。小さな会社内ならともかく一般に普及させるには労力が必要そうですね。
おわりに
PyTorch1.9からモデルをパッケージングするための機能 torch.package モジュールが追加されました。データサイエンティストとか研究者の方々にとっては他者による実験の再現を行うのに有用ですし、プロダクションに適用するエンジニアにとっては公式フォーマットでパッケージングされていることでCI/CDの仕組みに載せやすいというメリットがあります。得体の知れないzipファイルだとセキュリティチェックに引っかかることもありますので。
一方で、自由度が高い故に多少冗長なところもあります。依存ファイルの多いモデルのパッケージングは少し面倒かもしれません。ある程度制約を設けることでパッケージングの際の労力が減らせるといいかなと思いました。とはいえ、本機能はPythonモジュールをsdistとかbdistで作ったことがあるエンジニアなら簡単に扱えるでしょう。また、MLOps系ツールは世の中にたくさんありますが、PyTorch公式のフォーマットであれば正式対応してくれる可能性も期待できそうです。パッケージング処理のコードをわざわざ書かなくても、MLOpsツール側で自動でパッケージングしてくれると楽ですね。
本機能はまだβ版ということもあるので荒削りな部分はありますが、今後もユーザーの意見を取り入れながらブラッシュアップされていくといいですね。興味があれば是非使ってみましょう。
- 今回使った学習部分のソースコード(ipynb) alexnet_cifar10.ipynb