PyTorch 1.5 レシピ : TorchScript : 配備のための TorchScript (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 05/19/2020 (1.5.0)

* 本ページは、PyTorch 1.5 Recipes の以下のページを翻訳した上で適宜、補足説明したものです：

• TorchScript : TorchScript for Deployment

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

TorchScript : 配備のための TorchScript

このレシピでは、以下を学習します :

• TorchScript とは何か
• 訓練モデルを TorchScript 形式でどのようにエクスポートするか
• TorchScript モデルを C++ でどのようにロードして推論するか

要件

• PyTorch 1.5
• TorchVision 0.6.0
• libtorch 1.5
• C++ compiler

3 つの PyTorch コンポーネントをインストールするための手順は pytorch.org で利用可能です。C++ コンパイラは貴方のプラットフォームに依拠します。

TorchScript とは何か？

TorchScript は (nn.Module のサブクラス) PyTorch モデルの中間的な表現でそして C++ のような高パフォーマンス環境で実行可能です。それは PyTorch の高パフォーマンスなサブセットで、モデルの計算で実行時最適化を遂行する PyTorch JIT により消費されることを意味します。TorchScript は PyTorch モデルでスケールされた推論を行なうための推奨モデル形式です。より多くの情報については、PyTorch TorchScript へのイントロダクション チュートリアル、C++ で TorchScript モデルをロードする チュートリアルと 完全な TorchScript ドキュメント を見てください、それら総ては pytorch.org で利用可能です。

モデルをどのようにエクスポートするか

例として、事前訓練されたビジョン・モデルを取りましょう。TorchVision の事前訓練モデルの総ては TorchScript と互換です。

次の Python 3 コードをスクリプト、または REPL から実行します :

import torch
import torch.nn.functional as F
import torchvision.models as models

r18 = models.resnet18(pretrained=True)       # We now have an instance of the pretrained model
r18_scripted = torch.jit.script(r18)         # *** This is the TorchScript export
dummy_input = torch.rand(1, 3, 224, 224)     # We should run a quick test

2 つのモデルの等値性をサニティチェックしましょう :

unscripted_output = r18(dummy_input)         # Get the unscripted model's prediction...
scripted_output = r18_scripted(dummy_input)  # ...and do the same for the scripted version

unscripted_top5 = F.softmax(unscripted_output, dim=1).topk(5).indices
scripted_top5 = F.softmax(scripted_output, dim=1).topk(5).indices

print('Python model top 5 results:\n  {}'.format(unscripted_top5))
print('TorchScript model top 5 results:\n  {}'.format(scripted_top5))

モデルの両バージョンが同じ結果を与えることを見るはずです :

Python model top 5 results:
  tensor([[463, 600, 731, 899, 898]])
TorchScript model top 5 results:
  tensor([[463, 600, 731, 899, 898]])

そのチェックが確認されたら、続けてモデルをセーブします :

r18_scripted.save('r18_scripted.pt')

TorchScript モデルを C++ でロードする

次の C++ ファイルを作成してそれを ts-infer.cpp と名前付けます :

#include <torch/script.h>
#include <torch/nn/functional/activation.h>


int main(int argc, const char* argv[]) {
    if (argc != 2) {
        std::cerr << "usage: ts-infer <path-to-exported-model>\n";
        return -1;
    }

    std::cout << "Loading model...\n";

    // deserialize ScriptModule
    torch::jit::script::Module module;
    try {
        module = torch::jit::load(argv[1]);
    } catch (const c10::Error& e) {
        std::cerr << "Error loading model\n";
        std::cerr << e.msg_without_backtrace();
        return -1;
    }

    std::cout << "Model loaded successfully\n";

    torch::NoGradGuard no_grad; // ensures that autograd is off
    module.eval(); // turn off dropout and other training-time layers/functions

    // create an input "image"
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::rand({1, 3, 224, 224}));

    // execute model and package output as tensor
    at::Tensor output = module.forward(inputs).toTensor();

    namespace F = torch::nn::functional;
    at::Tensor output_sm = F::softmax(output, F::SoftmaxFuncOptions(1));
    std::tuple<at::Tensor, at::Tensor> top5_tensor = output_sm.topk(5);
    at::Tensor top5 = std::get<1>(top5_tensor);

    std::cout << top5[0] << "\n";

    std::cout << "\nDONE\n";
    return 0;
}

このプログラムは :

• コマンドラインで指定するモデルをロードする
• ダミー「画像」入力 tensor を作成する
• 入力上で推論を遂行する

また、このコードで TorchVison への依拠がないことにも気付くでしょう。TorchScript モデルのセーブされたバージョンは学習重みと計算グラフを持ちます - 他のものは必要ありません。

C++ 推論エンジンを構築して実行する

次の CMakeLists.txt ファイルを作成する :

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)

add_executable(ts-infer ts-infer.cpp)
target_link_libraries(ts-infer "${TORCH_LIBRARIES}")
set_property(TARGET ts-infer PROPERTY CXX_STANDARD 11)

プログラムを make する :

cmake -DCMAKE_PREFIX_PATH=<path to your libtorch installation>
make

今では、C++ で推論を実行して、結果を得ることを検証できます :

$ ./ts-infer r18_scripted.pt
Loading model...
Model loaded successfully
 418
 845
 111
 892
 644
[ CPULongType{5} ]

DONE

 
以上