PyTorch 1.4 Tutorials : PyTorch モデル配備 : TorchScript へのイントロダクション (翻訳/解説)
翻訳 : (株)クラスキャットセールスインフォメーション
作成日時 : 01/19/2020 (1.4.0)

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

Deploying PyTorch Models in Production : Introduction to TorchScript

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

PyTorch モデル配備 : TorchScript へのイントロダクション

このチュートリアルは TorchScript へのイントロダクションです、C++ のような高パフォーマンス環境内で実行可能な Python モデル (nn.Module のサブクラス) の中間表現です。

このチュートリアルでは以下をカバーします :

以下を含む、PyTorch のモデル・オーサリングの基本 :
- モジュール
- forward 関数を定義する
- モジュールをモジュールの階層に構成する
PyTorch モジュールを TorchScript、高パフォーマンス配備ランタイムに変換するための特定のメソッド
- 既存のモジュールを追跡する
- モジュールを直接コンパイルするためにスクリプティングを使用する
- 両者のアプローチをどのように構成するか
- TorchScript モジュールをセーブしてロードする

このチュートリアルを完了した後、C++ から TorchScript モデルを実際に呼び出すサンプルを調べる次の段階のチュートリアルを調べることに進むことを希望します。

import torch  # This is all you need to use both PyTorch and TorchScript!
print(torch.__version__)

1.4.0

PyTorch モデル・オーサリングの基本

単純な Module を定義することから始めましょう。Module は PyTorch の構成の基本ユニットです。それは以下を含みます :

コンストラクタ、これは発動のためにモジュールを準備します。
パラメータとサブモジュールのセット。これらはコンストラクタにより初期化されて発動の間モジュールにより使用できます。
forward 関数。これはモジュールが起動されたときに実行されるコードです。

小さいサンプルを調べましょう :

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()

    def forward(self, x, h):
        new_h = torch.tanh(x + h)
        return new_h, new_h

my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))

(tensor([[0.7108, 0.3658, 0.4292, 0.6288],
        [0.7710, 0.2918, 0.8513, 0.6062],
        [0.8313, 0.7276, 0.9428, 0.7750]]), tensor([[0.7108, 0.3658, 0.4292, 0.6288],
        [0.7710, 0.2918, 0.8513, 0.6062],
        [0.8313, 0.7276, 0.9428, 0.7750]]))

そこで私達は以下を行ないました :

torch.nn.Module をサブクラス化するクラスを作成しました。
コンストラクタを定義しました。コンストラクタは大きな働きはせず、単に super のためにコンストラクタを呼び出します。
forward 関数を定義しました、これは 2 つの入力を取りそして 2 つの出力を返します。forward 関数の実際の内容は実際には重要ではありませんが、それは fake RNN セルの類 – つまり – それは関数です、これはループ上で適用されます。

モジュールをインスタンス化して、x と y を作成します、これは単にランダム値の 3×4 行列です。それからセルを my_cell(x, h) で呼び起こします。これは forward 関数を順番に呼び出します。

もう少し興味深いことを行ないましょう :

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

MyCell(
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[0.6371, 0.5862, 0.0320, 0.8656],
        [0.7063, 0.1737, 0.5467, 0.4897],
        [0.6464, 0.4770, 0.7529, 0.2052]], grad_fn=<TanhBackward>), tensor([[0.6371, 0.5862, 0.0320, 0.8656],
        [0.7063, 0.1737, 0.5467, 0.4897],
        [0.6464, 0.4770, 0.7529, 0.2052]], grad_fn=<TanhBackward>))

モジュール MyCell を再定義しましたが、今回は self.linear を追加しました、そして forward 関数で self.linear を呼び出します。

ここで正確に何が起きているのでしょう？torch.nn.Linear は PyTorch 標準ライブラリからの Module です。丁度 MyCell のように、それは call シンタクスを使用して呼び出せます。Module の階層を構築しています。

Module 上の print は Module のサブクラス階層の視覚表現を与えます。私達のサンプルでは、Linear サブクラスとそのパラメータを見ることができます。

Module をこのように構成することにより、再利用可能なコンポーネントでモデルを簡潔に可読に作成できます。

出力上の grad_fn に気付いたかもしれません。これは autograd と呼ばれる自動微分の PyTorch のメソッドの詳細です。要するに、このシステムは潜在的に複雑なプログラムを通して導関数を計算することを可能にします。

この設計はモデル・オーサリングにおいて莫大な量の柔軟性を可能にします。

さて柔軟性と呼ばれるものを調べましょう :

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.dg = MyDecisionGate()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

MyCell(
  (dg): MyDecisionGate()
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.7138,  0.5495, -0.3841,  0.3377],
        [ 0.8317,  0.4166,  0.4498, -0.2328],
        [ 0.8537,  0.7811,  0.7466, -0.4516]], grad_fn=), tensor([[ 0.7138,  0.5495, -0.3841,  0.3377],
        [ 0.8317,  0.4166,  0.4498, -0.2328],
        [ 0.8537,  0.7811,  0.7466, -0.4516]], grad_fn=))

再度 MyCell クラスを再定義しましたが、ここでは MyDecisionGate を定義しました。このモジュールは制御フローを利用します。制御フローはループと if-ステートメントのようなものから成ります。

多くのフレームワークは完全なプログラム表現が与えられたときシンボリックな導関数を計算するアプローチを取ります。けれども、PyTorch では、勾配テープを使用します。演算をそれが発生するときに記録して、導関数を計算する間はそれらを後方に再生します。このように、フレームワークは言語の総ての構成物のために導関数を明示的に定義しなくてもかまいません。

autograd がどのように動作するか

TorchScript の基本

さて実行するサンプルを取り TorchScript をどのように適用できるかを見てみましょう。

要するに、TorchScript は PyTorch の柔軟で動的な性質の視点でさえも、モデルの定義を捕捉するツールを提供します。tracing と呼ぶものを調べることから始めましょう。

モジュールをトレースする

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

MyCell(
  original_name=MyCell
  (linear): Linear(original_name=Linear)
)

少し巻き戻して MyCell クラスの 2 番目のバージョンを取ります。前のように、それをインスタンス化しますが、今回は、torch.jit.trace を呼び出し、Module を渡し、そしてネットワークが見るかもしれないサンプル入力を渡します。

これは正確には何を成したのでしょう？それは Module を呼び出し、Module が実行されたときに発生した演算を記録し、そして torch.jit.ScriptModule のインスタンスを作成しました (TracedModule がそのインスタンスです)。

TorchScript は中間表現 (or IR) にあるその定義を記録します、一般には深層学習ではグラフとして参照されます。.graph プロパティでグラフを調べることができます :

print(traced_cell.graph)

graph(%self.1 : __torch__.torch.nn.modules.module.___torch_mangle_1.Module,
      %input : Float(3, 4),
      %h : Float(3, 4)):
  %19 : __torch__.torch.nn.modules.module.Module = prim::GetAttr[name="linear"](%self.1)
  %21 : Tensor = prim::CallMethod[name="forward"](%19, %input)
  %12 : int = prim::Constant[value=1]() # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %13 : Float(3, 4) = aten::add(%21, %h, %12) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %14 : Float(3, 4) = aten::tanh(%13) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %15 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%14, %14)
  return (%15)

けれども、これは非常に低位な表現でグラフに含まれる情報の殆どはエンドユーザには有用ではありません。代わりに、コードの Python-シンタクス解釈を与えるために .code プロパティを使用できます :

print(traced_cell.code)

def forward(self,
    input: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = torch.add((self.linear).forward(input, ), h, alpha=1)
  _1 = torch.tanh(_0)
  return (_1, _1)

それでは何故この総てを行なったのでしょう？幾つかの理由があります :

TorchScript はそれ自身のインタープリターで呼び出せます、これは基本的には制限された Python インタープリターです。このインタープリターは Global Interpreter Lock を獲得しませんので、非常に多くのリクエストが同じインスタンス上で同時に処理できます。
この形式はモデル全体をディスクにセーブしてそれを Python 以外の言語で書かれたサーバ内のような、もう一つの環境にロードすることを可能にします。
TorchScript はより効率的な実行を提供するためにコード上のコンパイラ最適化を行なうことができる表現を与えます。
TorchScript は個々の演算子よりもプログラムのより広いビューを要求する多くのバックエンド/デバイス・ランタイムと連動することを可能にします。

traced_cell の呼び出しは Python モジュールと同じ結果を生成することを見ることができます :

print(my_cell(x, h))
print(traced_cell(x, h))

(tensor([[ 0.6008, -0.3736,  0.7119,  0.1289],
        [-0.1579,  0.6616,  0.2627,  0.5846],
        [-0.0327,  0.7622,  0.4805,  0.1679]], grad_fn=<TanhBackward>), tensor([[ 0.6008, -0.3736,  0.7119,  0.1289],
        [-0.1579,  0.6616,  0.2627,  0.5846],
        [-0.0327,  0.7622,  0.4805,  0.1679]], grad_fn=))
(tensor([[ 0.6008, -0.3736,  0.7119,  0.1289],
        [-0.1579,  0.6616,  0.2627,  0.5846],
        [-0.0327,  0.7622,  0.4805,  0.1679]],
       grad_fn=), tensor([[ 0.6008, -0.3736,  0.7119,  0.1289],
        [-0.1579,  0.6616,  0.2627,  0.5846],
        [-0.0327,  0.7622,  0.4805,  0.1679]],
       grad_fn=<DifferentiableGraphBackward>))

モジュールを変換するためにスクリプティングを使用する

制御フローを積んだサブモジュールを持つ一つではなく、モジュールのバージョン 2 を使用したのには理由があります。それを今調べましょう :

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self, dg):
        super(MyCell, self).__init__()
        self.dg = dg
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell.code)

def forward(self,
    input: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = (self.dg).forward((self.linear).forward(input, ), )
  _1 = torch.tanh(torch.add(_0, h, alpha=1))
  return (_1, _1)

.code 出力を見ると、if-else 分岐はどこにも見つからないことを見ることができます！何故でしょう？tracing は正確に言ったことを行ないます、それは : コードを実行し、発生する演算を記録しそして正確にそれを行なう ScriptModule を構築します。不幸なことに、制御フローのようなものは消し去られます。

このモジュールを TorchScript でどのように忠実に表現できるのでしょう？スクリプトコンパイラを提供します、これはそれを TorchScript に変換するために Python ソースコードの直接的解析を行ないます。スクリプトコンパイラを使用して MyDecisionGate を変換しましょう :

scripted_gate = torch.jit.script(MyDecisionGate())

my_cell = MyCell(scripted_gate)
traced_cell = torch.jit.script(my_cell)
print(traced_cell.code)

def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = (self.dg).forward((self.linear).forward(x, ), )
  new_h = torch.tanh(torch.add(_0, h, alpha=1))
  return (new_h, new_h)

Hooray! We’ve now faithfully captured the behavior of our program in TorchScript. Let’s now try running the program:

# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell(x, h)

スクリプティングと tracing をミックスする

ある状況ではスクリプティングよりも tracing を使用することを必要とします (e.g. モジュールが TorchScript では現れて欲しくない定数 Python 値を元に作成される多くのアーキテクチャ決定を持つ (場合))。この場合、スクリプティングは tracing とともに組み合わせられます : torch.jit.script は traced モジュールのためにコードをインライン化できて、そして tracing は scripted モジュールのためにコードをインライン化できます。

最初のケースのサンプルは :

class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super(MyRNNLoop, self).__init__()
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))

    def forward(self, xs):
        h, y = torch.zeros(3, 4), torch.zeros(3, 4)
        for i in range(xs.size(0)):
            y, h = self.cell(xs[i], h)
        return y, h

rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)

def forward(self,
    xs: Tensor) -> Tuple[Tensor, Tensor]:
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y0 = y
  h0 = h
  for i in range(torch.size(xs, 0)):
    _0 = (self.cell).forward(torch.select(xs, 0, i), h0, )
    y1, h1, = _0
    y0, h0 = y1, h1
  return (y0, h0)

そして 2 番目のケースのサンプルは :

class WrapRNN(torch.nn.Module):
    def __init__(self):
        super(WrapRNN, self).__init__()
        self.loop = torch.jit.script(MyRNNLoop())

    def forward(self, xs):
        y, h = self.loop(xs)
        return torch.relu(y)

traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)

def forward(self,
    argument_1: Tensor) -> Tensor:
  _0, h, = (self.loop).forward(argument_1, )
  return torch.relu(h)

このように、スクリプティングと tracing は状況がそれらのいずれかを要求するとき一緒に使用できます。

モデルをセーブしてロードする

TorchScript モジュールをアーカイブ・フォーマットでディスクにセーブして (ディスクから) ロードするための API を提供します。このフォーマットはコード、パラメーター、属性そしてデバッグ情報を含みます、これはアーカイブが (完全に分離したプロセスにロード可能な) モデルの独立した表現であることを意味します。ラップされた RNN モジュールをセーブしてロードしましょう :

traced.save('wrapped_rnn.zip')

loaded = torch.jit.load('wrapped_rnn.zip')

print(loaded)
print(loaded.code)

RecursiveScriptModule(
  original_name=Module
  (loop): RecursiveScriptModule(
    original_name=MyRNNLoop
    (cell): RecursiveScriptModule(
      original_name=Module
      (dg): RecursiveScriptModule(original_name=MyDecisionGate)
      (linear): RecursiveScriptModule(original_name=Module)
    )
  )
)
def forward(self,
    argument_1: Tensor) -> Tensor:
  _0, h, = (self.loop).forward(argument_1, )
  return torch.relu(h)

見れるように、シリアライゼーションはモジュール階層と通して調べたコードを保存します。モデルはまた例えば、python-free な実行のために C++ にロードできます。

PyTorch 1.4 Tutorials : PyTorch モデル配備 : TorchScript へのイントロダクション