PyTorchの新しい推論モードについて

ひさしぶりのエントリー。今年度になってから家庭環境がガラッと変わってなかなかプライベートの一人の時間が取れず時間が空いてしまいました。

今回はPyTorch関連のメモ的な記事になります。最近仕事で古いPyTorchで作られたプロダクトを最新のPyTorchにアップグレードする対応をしていたのですが、バージョン1.4からいきなり最新バージョンと1.9に上げたのでちょっと大変でした。ここでは1.9から入った新機能の一つを紹介したいと思います。地味ではありますが、意外と気になる人が多そうな機能を調べてみました。

* PyTorch 1.9 Release, including torch.linalg and Mobile Interpreter | PyTorch

Inference Mode

Inference Mode API allows significant speed-up for inference workloads while remaining safe and ensuring no incorrect gradients can ever be computed. It offers the best possible performance when no autograd is required.

Performance Optimization and Toolingの項目に、Inference Mode APIなるものがBeta版として追加されています。文字通り、推論モード用のAPIということなんですけど、PyTorchを以前から使っていた人は一見困惑するのではないかと思います。

まず整理としてこれまで推論時にどのような対応をしていたかというと、

model.eval()  ## torch.nn.Module.eval
with torch.no_grad():  ## disable autograd
    model(data)  # forward

model.eval() ## torch.nn.Module.eval

with torch.no_grad(): ## disable autograd

model(data) # forward

意味としては、評価モード(Dropouts Layers、BatchNorm Layersをスキップ)に切り替えて、自動微分を無効(勾配計算用パラメータを保存しないNoGrad Mode)にしてから実行することで不要な処理、無駄なメモリ消費を抑えて推論を実行することができます。torch.no_grad() は以下のようにデコレータとしても指定できます。

@torch.no_grad()
def doubler(x):
    return x * 2
z = doubler(x)
print(z.requires_grad)  ## False

@torch.no_grad()

def doubler(x):

return x * 2

z = doubler(x)

print(z.requires_grad) ## False

1.9で追加された新しい推論モード torch.inference_mode() も従来の torch.no_grad() と同様にコンテキストマネージャーとして提供されておりPythonレベルでの実装はこれまでとそう変わらないようにデザインされています。

x = torch.ones(1, 2, 3, requires_grad=True)
with torch.inference_mode():
    y = x * x
print(y.requires_grad)  ## False

@torch.inference_mode()
def func(x):
    return x * x
out = func(x)
print(out.requires_grad)  ## False

x = torch.ones(1, 2, 3, requires_grad=True)

with torch.inference_mode():

y = x * x

print(y.requires_grad) ## False

@torch.inference_mode()

def func(x):

return x * x

out = func(x)

print(out.requires_grad) ## False

inference_mode を使った場合でも requires_grad=False になります。重要なのは torch.no_grad() と何が違うかということなので実装を見てみます。

torch.inference_mode は前述の通りコンテキストマネージャーなのでPythonのclassとして実装されています。
* pytorch/grad_mode.py at master · pytorch/pytorch · GitHub

class inference_mode(_DecoratorContextManager):
... 省略
    def __init__(self, mode=True):
        if not torch._jit_internal.is_scripting():
            super().__init__()
        # Holds a python binding to a RAII guard that can enable or disable
        # inference mode
        self._inference_mode_raii_guard = None
        self.mode = mode

class inference_mode(_DecoratorContextManager):

... 省略

def __init__(self, mode=True):

if not torch._jit_internal.is_scripting():

super().__init__()

# Holds a python binding to a RAII guard that can enable or disable

# inference mode

self._inference_mode_raii_guard = None

self.mode = mode

RAIIなる単語が出てきていますが、Resource Acquisition Is Initializationの略で，リソースの確保(Acquisition)と解放を変数の初期化(Initialization)と破棄に紐付けるというテクニックです。C++を書くプログラマには馴染みは深いかと思いますし、僕も学生時代にEffective C++を読んで学びました。最近流行りのRustだとRAIIが強制されるのでメモリ安全性を備えていると言われてますね。

Python実装内では単にモード切替フラグの管理だけで実装の本体は低レイヤー(C++)の方にあります。
* pytorch/TensorImpl.h at master · pytorch/pytorch · GitHub

ちなみにC10というのは少し前からPyTorchのバックエンド実装に用いられているコアライブラリを指します。ざっくり言うとCaffe2内部でも使っているATen(テンソル演算用C++モジュール)を統合したものです。C10という名前はCore TENsor Libraryからきているみたいな話もありますけど公式の由来なのかはわかりません。あまり興味もないですが。

細かい実装はATenの詳細に入るので省略しますが、ざっくり言うとInference Modeの場合は、そのブロック内でのみ使われる推論専用のTensorが用意され(RAII)、推論計算時にはInplace(破壊的)に値が更新されます。

また、この推論用に用意されるTensorには一部制限があり、Inference Modeの外側ではTensor内の値は変更不可、変更しようとするとRuntimeErrorとなります。プロダクションでの推論処理用途なら特に問題は無いかと思います。また、Inference Modeで作られたTensorに対してView操作をしても追跡されません。

x = torch.ones(1, 2, 3, requires_grad=True)
with torch.inference_mode():
    y = x * x
y[0][0][1] = 2
RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed.You can make a clone to get a normal tensor before doing inplace update.See https://github.com/pytorch/rfcs/pull/17 for more details.

x = torch.ones(1, 2, 3, requires_grad=True)

with torch.inference_mode():

y = x * x

y[0][0][1] = 2

RuntimeError: Inplace update to inference tensor outside InferenceMode is not allowed.You can make a clone to get a normal tensor before doing inplace update.See https://github.com/pytorch/rfcs/pull/17 for more details.

また、Inference ModeのTensorは以下のような内部動作となります。

torch.no_grad() と同様に自動微分での勾配計算は無効になり、grad_fn を始め計算グラフの情報は常に記録されない
Version Counter(torch.Tensor._version)は保存されず、全てのオペレーションが追跡されない

Version Counterというのは自動微分するために内部で必要となるデータバージョニング用のカウンター変数ですが推論用には不要なので保存されません。詳細については自動微分の実装に関わってくるのでここでは省略しますが、気になる方はATenの ADInplaceOrView あたりから読んでいくといいかと思います。ちなみにカウンターの値自体は以下のように簡単に確認できます。

v = values.add_(1).add_(1).sub_(1).sub_(1)
print(v._version)  ## => 4 (オペレーション毎にカウントアップされるデータバージョニング用変数)

## 推論モード
with torch.inference_mode():
    y = x * x
print(y._version)
## Inference Modeではversion counterは保存されない
Traceback (most recent call last):
  File "./test.py", line 14, in <module>
    print(y._version)
RuntimeError: Inference tensors do not track version counter.

v = values.add_(1).add_(1).sub_(1).sub_(1)

print(v._version) ## => 4 (オペレーション毎にカウントアップされるデータバージョニング用変数)

## 推論モード

with torch.inference_mode():

y = x * x

print(y._version)

## Inference Modeではversion counterは保存されない

Traceback (most recent call last):

File "./test.py", line 14, in <module>

print(y._version)

RuntimeError: Inference tensors do not track version counter.

torch.no_grad() ではVersion Counterは保存していますし、Tensor内の値も後で変更可能です。Inference Modeは低レイヤー実装内でメモリ使用の無駄を省く処理がいろいろ施されているようでした。公式にも the extreme version of no-grad mode. としているので内部実装ではそこそこ差はあるようです。

おわりに

PyTorch 1.9から新しい推論モード torch.inference_mode() が追加されました。Pythonレベルだと、torch.no_grad() と同様にコンテキストマネージャーまたはデコレータとして利用することができるため書き換えは容易です。実体としては torch.no_grad() の単なるエイリアスというわけではなく、メモリ効率が改善された新たな推論処理特化の仕組みとなっています。実用上の不都合についてはまだプロダクションでは使っていないので性能面での定量評価もできていませんが、環境が整ったらより詳細を調査してみようと思います。

公式でも It is recommended that you try out inference mode in the parts of your code that do not require autograd tracking って言ってるのでぜひ試してみてください。

PyTorchの新しい推論モードについて

Inference Mode

おわりに

あわせて読む:

コメントを残す