サーバーレスGPU推論（Cerebrium）

こんにちは！Fusic 機械学習チームのトシュコフです。

機械学習モデルの開発から運用までなんでもしています。もし、機械学習で困っていることがあれば、気軽にお問い合わせください。

---

本記事ではCerebriumという新しいサーバーレスGPUサービスを紹介します。

注意：まだベータであるため、不安定が多くバグもありますが、Slackチャンネルでの対応が非常によく、問題を経験した時に是非連絡してみてください。

概要

AIモデルをデプロイする時、GPUを利用することによって推論時間を非常に押せえる事ができます。リアルタイム推論をする時に推論時間が非常に大事になりますので、GPUで推論することが望ましいです。
ただし、クラウドのGPUインスタンスで推論する時にいくつかの問題が発生します：

1. 1つのインスタンスをずっと借りる事で無駄の費用が発生します。
2. TorchServeのようなツールがありますが、設定が難しい場合があります。
3. リクエスト数によってGPUインスタンスを自動的に合わせるができますが、その場合はトラフィックルーティングが難しくなります。

そのために、AWS Lambdaのようなサーバーレス推論環境が最も管理しやすいですが、GPU推論ができないため、SageMaker Serverles Inferenceのようなサービスを使う必要があります。

今回はCerebriumという費用が低く、かなり使いやすいサーバレスGPUサービスを試してみました。

デプロイ可能モデル

デプロイ可能のモデルは8種類あります：

1. PyTorch
2. ONNX
3. Scikit
4. XGBoost
5. TensorFlow
6. Spacy
7. HuggingFace Transformers
8. カスタムPythonコード

Cerebriumの特徴として、たった1行でモデルをデプロイできる事です。

endpoint = deploy((model_type.TORCH, "torch_model.pkl"), "torch-model-cp", "<API_KEY>")

カスタムPythonコードのデプロイ方法が他のサービスと同様で、Pythonコード＋requirements.txtを使用します。

PyTorchモデルのデプロイ方法について

PyTorch 2.0が出たのに関わらず、依存モジュールも含まれるモデルを保存する方法がまだ難しいです。モデルの重みを保存し、モデルの定義を依存モジュールとモデルの定義ファイルから読み込む事になっています。

その理由で、この問題のCerebriumの解決方法が非常に面白いかと思いました。Pickleが依存物を保存できないため、CloudPickleという保存可能な物がより広いライブラリーが使用されています。

class NeuralNet(nn.Module):
    def __init__(self, num_classes):
        super(NeuralNet, self).__init__()
        self.conv_layer = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
	...
        self.fc = nn.Linear(128, num_classes)

    def forward(self, x):
        out = self.conv_layer(x)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

model = NeuralNet(10)

## Cloudpickleで保存
import cloudpickle
with open("torch_model.pkl", "wb") as f:
    cloudpickle.dump(model, f)

deploy((model_type.TORCH, "torch_model.pkl"), "torch-model-cp", "<API_KEY>")

ただし、他のモジュールの依存があった場合、CloudPickleでも対応できない事が多いです。
そのためにONNX変換を行い、それをアップロードすることをお勧めします。

ONNXモデルのデプロイ方法について

ONNXとはPyTorchやTensorFlowのモデルもエキスポートできるフォーマットです。

注意：PyTorchの全てのモデルがエキスポートできませんが、カスタム関数を実装する必要があるモデルもあります。

PyTorchモデルをエキスポートすることは以下のコマンドで行えます：

input_names = ["input"]
output_names = ["output"]
torch.onnx.export(
    model,
    images,
    "pytorch.onnx",
    verbose=True,
    input_names=input_names,
    output_names=output_names,
)

その次に、Cerebriumにデプロイすることは以下の1つのコマンドで行えます。

from cerebrium import deploy, model_type
output_flow = deploy((model_type.ONNX, "pytorch.onnx"),"onnx-pytorch", "<API_KEY>")

ONNXモデルを推論する時にローカルのGPUで推論する時よりも早いレスポンスが多々あって、非常に驚きました。気になったので、Cerebriumの社長と会話をしてみました。アップロードされるモデルが1つのONNXファイルという決まっている形式であるため、アップロードする際に最先端の最適化技術を加える事ができるそうです。大変素晴らしいアップローチで、エンドユーザーの負担を減らしています。

前処理

最近追加された機能ですが、ファイルアップロードもできるようになっています。

入力が画像の場合、APIリクエストする時に全ての画素をリストで送る場合、レスポンス時間に大きく影響するため、カスタムな前処理として画像をバイトとして読込、サーバーで行列に変換を行った方が良いです。

def pre_process(data, files):
    import numpy as np
    from PIL import Image
    labelled = {d: Image.open(f) for d, f in zip(data, files)}
    return labelled

from cerebrium import Conduit, model_type
conduit = Conduit("onnx-pytorch", '<API_KEY>', [(model_type.ONNX, "pytorch.onnx", {"pre": pre_process"})])
conduit.load('./')
conduit.run(data, [files]) # files should be binary objects - for images it would the return of Image.open(f)
conduit.deploy()

全ての前処理が対応されていませんが、現在は以下のライブラリーが使えます：

1. numpy
2. scikit-learn
3. torch
4. pandas
5. Pillow
6. transformers

まとめ

この記事では、Cerebriumという新たなサーバレスGPUサービスを紹介しました。

エンドユーザとして非常に使いやすく、費用が低いサービスです。特にONNXモデルをアップロードすると最もいい結果を感じました。Cold-Startが早く、推論速度も非常に早いサービスで、是非試してみてください。

---

最後に宣伝になりますが、機械学習でビジネスの成長を加速するために、Fusicの機械学習チームがお手伝いしています。機械学習のPoCから運用まで、すべての場面でサポートした実績があります。もし、困っている方がいましたら、ぜひFusicにご相談ください。

お問い合わせ