PyTorch推論

PyTorchモデルの推論

PyTorchについて学び、PyTorchモデルで推論を実行する方法を理解してください。

PyTorchは、理解しやすく柔軟なAPI、特に自然言語処理（NLP）ドメインで利用可能な多数の既製モデル、およびドメイン固有のライブラリにより、深層学習の分野をリードしています。

PyTorchで構築されたモデルを使用しようとする開発者とアプリケーションのエコシステムが拡大しており、この記事ではPyTorchモデルの推論を簡単に紹介します。PyTorchモデルの推論を実行するには、いくつかの異なる方法があります。以下に列挙します。

この記事は、PyTorchモデルのトレーニング方法ではなく、PyTorchモデルで推論を実行する方法についての情報を探していることを前提としています。

目次

PyTorchの概要

PyTorchの核心はnn.Moduleで、これは深層学習モデル全体または単一のレイヤーを表すクラスです。モジュールは組み合わせたり拡張したりしてモデルを構築できます。独自のモジュールを作成するには、モデル入力とモデルのトレーニング済み重みに基づいて出力を計算するforward関数を実装します。独自のPyTorchモデルを作成している場合は、おそらくそれをトレーニングもしているでしょう。あるいは、PyTorch自体またはHuggingFaceなどの他のライブラリから事前トレーニング済みモデルを使用できます。

PyTorch自体を使用して画像処理モデルをコーディングする場合：

import torch
import torch.nn as nn
import torchvision.transforms as T
from torchvision.models import resnet18, ResNet18_Weights


class Predictor(nn.Module):

    def __init__(self):
        super().__init__()
        weights = ResNet18_Weights.DEFAULT
        self.resnet18 = resnet18(weights=weights, progress=False).eval()
        self.transforms = weights.transforms()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            x = self.transforms(x)
            y_pred = self.resnet18(x)
            return y_pred.argmax(dim=1)

HuggingFaceライブラリを使用して言語モデルを作成する場合：

model_name = "bert-large-uncased-whole-word-masking-finetuned-squad"

tokenizer = transformers.BertTokenizer.from_pretrained(model_name)
model = transformers.BertForQuestionAnswering.from_pretrained(model_name)

トレーニング済みモデルを作成またはインポートしたら、推論を実行するためにそれをどのように実行しますか？以下では、PyTorchで推論を実行するために使用できるいくつかのアプローチについて説明します。

推論オプション

ネイティブPyTorchでの推論

パフォーマンスやサイズに敏感ではなく、Python実行可能ファイルとライブラリを含む環境で実行している場合は、ネイティブPyTorchでアプリケーションを実行できます。

トレーニング済みモデルを取得したら、推論のためにモデルを保存および読み込むために、あなた（またはデータサイエンスチーム）が使用できる2つの方法があります：

1. モデル全体を保存および読み込む

   # モデル全体をPATHに保存
   torch.save(model, PATH)
   
   # PATHからモデルを読み込み、推論のためにevalモードに設定
   model = torch.load(PATH)
   model.eval()

2. モデルのパラメータを保存し、モデルを再宣言してパラメータを読み込む

   # モデルパラメータを保存
   torch.save(model.state_dict(), PATH)
   
   # モデルを再宣言し、保存されたパラメータを読み込む
   model = TheModel(...)
   model.load_state_dict(torch.load(PATH))
   model.eval()

これらの方法のどちらを使用するかは、構成によって異なります。モデル全体を保存および読み込むことは、モデルを再宣言する必要がなく、モデルコード自体にアクセスする必要がないことを意味します。しかし、トレードオフとして、保存環境と読み込み環境の両方が、利用可能なクラス、メソッド、パラメータの点で一致する必要があります（これらは直接シリアル化および逆シリアル化されるため）。

元のモデルコードにアクセスできる限り、モデルのトレーニング済みパラメータ（状態辞書、またはstate_dict）を保存することは、最初のアプローチよりも柔軟です。

ネイティブPyTorchを使用してモデルで推論を実行したくない主な理由は2つあります。1つ目は、PythonランタイムとPyTorchライブラリおよび関連する依存関係を含む環境で実行する必要があることです。これらは合計で数ギガバイトのファイルになります。携帯電話、Webブラウザ、または専用ハードウェアなどの環境で実行したい場合、ネイティブPyTorchでPyTorch推論を実行することは機能しません。2つ目はパフォーマンスです。そのままの状態では、PyTorchモデルはアプリケーションに必要なパフォーマンスを提供しない場合があります。

TorchScriptでの推論

PyTorchやその他のPythonライブラリをインストールできない、より制約のある環境で実行している場合、TorchScriptに変換されたPyTorchモデルで推論を実行するオプションがあります。TorchScriptはPythonのサブセットで、非Python環境で読み込んで実行できるシリアル化可能なモデルを作成できます。

# TorchScriptにエクスポート
script = torch.jit.script(model, example)

# スクリプト化されたモデルを保存
script.save(PATH)

# スクリプト化されたモデルを読み込む
model = torch.jit.load(PATH)
model.eval()

#include <torch/script.h>

...

  torch::jit::script::Module module;
  try {
    // ScriptModuleを逆シリアル化
    module = torch::jit::load(PATH);
  }
  catch (const c10::Error& e) {
    ...
  }

...

TorchScriptアプローチを使用してPyTorchモデルで推論を実行するために環境にPythonランタイムが必要ありませんが、libtorchバイナリをインストールする必要があり、これらは環境にとって大きすぎる場合があります。また、アプリケーションに必要なパフォーマンスが得られない場合もあります。

ONNX Runtimeでの推論

パフォーマンスと移植性が最重要である場合、ONNX Runtimeを使用してPyTorchモデルの推論を実行できます。ONNX Runtimeを使用すると、レイテンシとメモリを削減し、スループットを向上させることができます。また、ONNX Runtimeに付属の言語バインディングとライブラリを使用して、クラウド、エッジ、Web、またはモバイルでモデルを実行できます。

最初のステップは、PyTorch ONNXエクスポーターを使用してPyTorchモデルをONNX形式にエクスポートすることです。

# サンプルデータを指定
example = ...

# モデルをONNX形式にエクスポート
torch.onnx.export(model, PATH, example)

ONNX形式にエクスポートしたら、オプションでNetronビューアーでモデルを表示して、モデルグラフと入力および出力ノード名と形状、および可変サイズの入力と出力（動的軸）を持つノードを理解できます。

その後、選択した環境でONNXモデルを実行できます。ONNX RuntimeエンジンはC++で実装されており、C++、Python、C#、Java、Javascript、Julia、RubyのAPIがあります。ONNX Runtimeは、Linux、Mac、Windows、iOS、Androidでモデルを実行できます。たとえば、次のコードスニペットは、C++推論アプリケーションのスケルトンを示しています。

  // ONNX Runtimeセッションを割り当て
  auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
  Ort::Env env;
  Ort::Session session{env, ORT_TSTR("model.onnx"), Ort::SessionOptions{nullptr}};

  // モデル入力を割り当て：形状とサイズを入力
  std::array<float, ...> input{};
  std::array<int64_t, ...> input_shape{...};
  Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input.data(), input.size(), input_shape.data(), input_shape.size());
  const char* input_names[] = {...};

  // モデル出力を割り当て：形状とサイズを入力
  std::array<float, ...> output{};
  std::array<int64_t, ...> output_shape{...};
  Ort::Value output_tensor = Ort::Value::CreateTensor<float>(memory_info, output.data(), output.size(), output_shape.data(), output_shape.size());
  const char* output_names[] = {...};

  // モデルを実行
  session_.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, &output_tensor, 1);

そのままの状態で、ONNX RuntimeはONNXグラフに一連の最適化を適用し、可能な場合はノードを結合し、定数値を因数分解します（定数フォールディング）。ONNX Runtimeは、Execution Providerインターフェースを介して、ターゲットとするハードウェアプラットフォームに応じて、CUDA、TensorRT、OpenVINO、CoreML、NNAPIなど、多数のハードウェアアクセラレータと統合されます。

モデルを量子化することで、ONNXモデルのパフォーマンスをさらに向上させることができます。

アプリケーションがモバイルやエッジなどの制約のある環境で実行されている場合、アプリケーションが実行するモデルまたはモデルセットに基づいて、縮小サイズのランタイムをビルドできます。

選択した言語と環境で開始するには、ONNX Runtimeの入門を参照してください。

例

モデルをONNXにエクスポート

• torch.onnxによる基本的なPyTorchエクスポート
• ONNX Runtimeによる超解像
• カスタム演算子を使用してPyTorchモデルをエクスポート

ONNX Runtimeによる高速化推論

• CPUでBERTモデルを高速化

• GPUでBERTモデルを高速化

• 量子化による縮小サイズBERTモデルの高速化

• CPUでGPT2を高速化

RESNET-50

• CPU、GPU、OpenVINOでRESNET-50モデルを高速化