AWS Neuron SDK を使ってローカルでPytorchの自作モデルをコンパイルしてみた

AWS Neuron SDKとは？

AWSでは、AWS Inferentiaという機械推論向けのチップを搭載したEC2のインスタンスが提供されています。
AWS Neuron SDKはAWS Inferentia上で推論を実施するためのSDKとなっています。
tensorflowやpytorchで作られたモデルを専用でコンパイルすることによって、AWS Inferentia上で実行できるモデルに変換することが可能です。

チュートリアルはどれもAWS上でインスタンスを立ててのコンパイルとなっていたのですが、ローカルのマシンで開発している身としてはちょっと手間のかかる内容だったので、今回ローカルでコンパイルしてみました。
また、ResNet50そのままを利用したものばかりだったので、自作したモデルをコンパイルしてみました。

実施環境

• Windows 10 WSL2 Ubuntu 18.04.2 LTS
• Python 3.7.11
  • Tensorflowの1系をライブラリの依存関係に要求されており、1系が対応していない3.8以降は利用できません(個人的にはハマりどころでした)。

必要ライブラリのインストール

まずはNeuron SDKのリポジトリを追加。

$ pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

その後ライブラリのインストール。

$ pip install torch-neuron=1.8.1.* torchvision==0.9.1 neuron-cc[tensorflow]

チュートリアルではバージョンを明言されていないことが多かったのですが、素直にpipにお任せすると、互換性のないtorch1.8.1とtorchvisionの0.10系がインストールされて、うまく行かないので注意です。

利用するライブラリのバージョンの組み合わせはこちら。

PyTorch	torch-neuron	torchvision	備考
1.8.1	1.8.1.*	0.9.1	2021/9時点での torch-neuron最新版
1.7.1	1.7.1.*	0.8.2	2021/9時点での AWS Deep Learning Containers 対応してる中での最新版

コンテナ内で利用する予定がない場合は、PyTorch1.8.1対応版、ある場合は1.7.1対応版といったところでしょうか。

ちなみにこの場合でもnumpyのバージョンが依存関係と合わない、と怒られますが、動作したので今回は無視しています。

ERROR: neuron-cc 1.6.13.0+9f61b2f75 has requirement numpy<=1.20.0,>=1.13.3, but you'll have numpy 1.21.2 which is incompatible.

コンパイルしてみる

今回のコンパイルで利用したモデルは、個人的に作成していたResNet18を利用した、イラスト/写真分類モデルです。
そのモデルのクラスを持ってきて、以下のようなスクリプトを作成しました。

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

import torch_neuron
from torchvision import models

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        
        self.conv = models.resnet18(pretrained=True)
        self.fc = nn.Linear(1000, 2)
        
    def lossfun(self, y, t):
        return F.cross_entropy(y, t)

    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=0.01)

    def forward(self, x):
        x = self.conv(x)
        x = self.fc(x)

        return x

image = torch.zeros([1, 3, 256, 256], dtype=torch.float32)

model = Net()
model.load_state_dict(torch.load('resources/2d-recog.pt', map_location=torch.device('cpu')))
model.eval()

model_neuron = torch.neuron.trace(model, example_inputs=[image])

model_neuron.save("2d-recog_neuron.pt")

このスクリプトを実行すると、無事コンパイルが成功しました。

Compiler status PASS
INFO:Neuron:Number of arithmetic operators (post-compilation) before = 72, compiled = 72, percent compiled = 100.0%
INFO:Neuron:The neuron partitioner created 1 sub-graphs
INFO:Neuron:Neuron successfully compiled 1 sub-graphs, Total fused subgraphs = 1, Percent of model sub-graphs successfully compiled = 100.0%
INFO:Neuron:Compiled these operators (and operator counts) to Neuron:
INFO:Neuron: => aten::_convolution: 20
INFO:Neuron: => aten::adaptive_avg_pool2d: 1
INFO:Neuron: => aten::add: 8
INFO:Neuron: => aten::addmm: 2
INFO:Neuron: => aten::batch_norm: 20
INFO:Neuron: => aten::flatten: 1
INFO:Neuron: => aten::max_pool2d: 1
INFO:Neuron: => aten::relu: 17
INFO:Neuron: => aten::t: 2

今回のスクリプトの例でいうと、2d-recog_neuron.ptも同一ディレクトリに生成されています。

こちらを利用することで、inferentiaチップを利用した推論ができるようになります。
ただし、ローカルで実行しても推論は行われず空撃ちとなる様子でした。

/home/xxxx/workspace/inferentia-test/.neuron/lib/python3.7/site-packages/torch/nn/modules/module.py:889: UserWarning: Neuron runtime cannot be initialized; falling back to CPU execution
Tensor output are ** NOT CALCULATED ** during CPU execution and only indicate tensor shape (Triggered internally at  /opt/workspace/KaenaPyTorchRuntime/neuron_op/neuron_op_impl.cpp:123.)

推論についても、チュートリアルがAWSインスタンスありきとなっていることが多いので、ローカルで作成したコンテナイメージをECS上にデプロイすることで実施できないか検証中です。
近日公開する予定ですので、続報をお待ちください。

参考文献

AWS Neuron SDK - github (Sampleのresnet50.ipynb)
AWS Neuron Documentation