【モデルデプロイ】TorchScript解釈(4)：Torch jitにおけるエイリアス解析

トーチスクリプトの解釈 (4): トーチ jit におけるエイリアス解析 - Zhihu (zhihu.com)

トーチスクリプトシリーズは引き続き更新されますので、皆さん長らくお待ちください。前回の研究では、torch jit の基本概念をいくつか習得し、Python で書かれたモデルを torchscript と ONNX に変換する方法を学び、いくつかの使いやすいツールを使用してパスを生成し、モデルを最適化することができました。

OpenMMLab：TorchScriptの解釈（１）：初めてTorchScriptを知る

OpenMMLab：TorchScript の解釈 (2)：Torch jit トレーサによる分析の実装 82 同意 · 5 コメント記事をアップロード中...再アップロードキャンセル

OpenMMLab：TorchScript の解釈 (3): jit のサブグラフ rewriter43 同意 · 0 コメント記事をアップロード中...再アップロードキャンセル

一部のリーダーは、最適化のニーズを満たすために、より複雑で強力なパスを作成できますが、より複雑なコードは、データハザードなどの潜在的なリスクが増えることを意味します。エイリアス分析は、リスクを回避し、より安全なパスを作成するのに役立つツールです。今日は、それについて一緒に理解していきます。

エイリアス解析とは

Torch jit には、さまざまな最適化を完了するのに役立つ組み込みパスが多数あり、ユーザーは特定の目的を達成するために独自のパスを定義することもできます。この柔軟性はモデルを最適化するのに便利ですが、制限がないわけではありません。例として、次のコードと対応するビジュアルイメージを取り上げます。

def forward(self, x, y): 
    x = x + 1 
    x.add_(x) 
    return x + y 
 
# graph(%self : __torch__.TestModel, 
#       %x.1 : Tensor, 
#       %y.1 : Tensor): 
#   %4 : int = prim::Constant[value=1]() # create_model.py:19:16 
#   %x0.1 : Tensor = aten::add(%x.1, %4, %4) # create_model.py:19:12 
#   %8 : Tensor = aten::add_(%x0.1, %x0.1, %4) # create_model.py:20:8 
#   %11 : Tensor = aten::add(%x0.1, %y.1, %4) # create_model.py:21:15 
#   return (%11)

ニューラルネットワークは計算グラフ Graph を形成します。原則として、Graph 内のノードがどのパスを通じても出力ノードに到達できない場合、これは無駄なノードであるため、最適化 (削除) できます。たとえば、上の図のノードです add_ 。この最適化は、デッドコードの除去 (DeadCodeElimination) と呼ばれることがよくあります。add_ コードから、このノードは x の値を更新するインプレース操作であり、ノードが削除されると誤った計算結果が生じることがわかります。

上記のコードの 11 行目からわかるように、add_ is の出力は、実際にその input と同じメモリ空間を共有していること%8がわかっている場合、最適化の正確性を確保するためにこのノードの削除を回避します。これがエイリアス分析の役割です。%8%x0.1

次のコードでそれを確認できます。

#include <torch/csrc/jit/ir/alias_analysis.h> 
#include <torch/script.h> 
 
int main(int argc, char* argv[]) { 
  auto model = torch::jit::load(argv[1]);                // 读取模型 
  auto graph = model.get_method("forward").graph();      // 提取计算图 
  torch::jit::AliasDb aliasdb(graph);                    // 创建AliasDb对象 
  aliasdb.dump();                                        // 可视化分析结果 
  return 0; 
}

このプログラムは、PyTorch が提供するエイリアス解析ツール AliasDb を利用して、入力モデルの順関数のグラフを解析し、解析結果を可視化します。今モデルを入力したところ、結果は次のようになります。

===1. GRAPH=== 
graph(%self : __torch__.TestModel, 
      %x.1 : Tensor, 
      %y.1 : Tensor): 
  %4 : int = prim::Constant[value=1]() # create_model.py:19:16 
  %x0.1 : Tensor = aten::add(%x.1, %4, %4) # create_model.py:19:12 
  %8 : Tensor = aten::add_(%x0.1, %x0.1, %4) # create_model.py:20:8 
  %11 : Tensor = aten::add(%x0.1, %y.1, %4) # create_model.py:21:15 
  return (%11) 
 
===2. ALIAS DB=== 
%x.1 points to: WILDCARD for type Tensor 
%y.1 points to: WILDCARD for type Tensor 
%8 points to: %x0.1 
%self points to: WILDCARD for type __torch__.TestModel 
 
===3. Writes=== 
%8 : Tensor = aten::add_(%x0.1, %x0.1, %4) # create_model.py:20:8 
  %x0.1,

ALIAS DB: の下にこのような行があることがわかります %8 points to: %x0.1。%8 このツールを使用すると、実際に参照された %x0.1 値が %x0.1 ネットワーク出力の計算に関与しているため、 %8 計算を削除すべきではないことがわかります。

AliasDb

AliasDb は、PyTorch が提供するエイリアス分析ツールです。AliasDb の助けを借りて、計算グラフ内の各データノードの関係を分析し、誤った最適化の可能性を回避できます。

メモリDAG

MemoryDAG はストレージグラフオブジェクトであり、データ間の依存関係を維持するために AliasDb によって使用されます。ソースコードによると https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/ir/alias_analysis.h 。AliasDb は計算グラフ Graph を受け取り、ストレージグラフ MemoryDAG を作成します。次の表に示すように、この MemoryDAG と Graph の間には対応関係があります。

	グラフ	メモリDAG
写真	Graph オブジェクトは計算グラフの構造を維持し、計算中のデータフローを記述します。	MemoryDAG オブジェクトはストレージグラフ構造を維持し、要素間の参照 (ポインティング) 関係を記述します。
ノード	ノードはそれぞれの独立した計算を表し、入力と出力は Value オブジェクトです	要素は、他の要素を参照するかどうかを含む、1 つ (または複数) の Value オブジェクトのデータストレージ情報を表します。
側	Use オブジェクトは、値がどのノードで使用されるかを示します。	MemoryLocations オブジェクトは、その要素が参照できる他の要素を示します。

注意すべき点がいくつかあります。

要素は値を指すだけでなく、コンテナクラスやワイルドカードなどを指すこともあります。
Element と Value は必ずしも 1 対 1 に対応しているわけではありません。たとえば、分岐構造では、 if condition: val=A else: val=B Element が A と B のいずれかを指すことができます。
上記の理由に基づいて、MemoryLocations が指す要素は 1 つの可能性を表すだけです。

AliasDb はグラフを受け取ると、各ノードの FunctionSchema で提供された情報に従って対応する MemoryDAG オブジェクトを構築し、その後の分析を容易にします。

FunctionSchema が何であるかを覚えていない場合は、 Torch jit トレーサ実装分析で基本的な知識を復習できます。

view FunctionSchema は次のとおりです。

view(Tensor(a) self, int[] size) -> Tensor(a)

最初のパラメーター self と出力にマークがあることがわかります (a) 。これは、出力がパラメーター self のエイリアスである可能性があることを示しています。

以下は、計算グラフ Graph と、対応するストレージグラフ MemoryDAG の視覚化を表すコードです。

@torch.jit.script 
def foo(a : Tensor, b : Tensor): 
  c = 2 * b 
  a += 1 
  if a.max() > 4: 
    r = a[0] 
  else: 
    r = b[0] 
  return c, r

次のコードに示すように、 AliasDb を使用する と、出力以外のどの Value が Node によって読み書きされるかを簡単にクエリできます。これは、将来のデータハザードにどのように対処するかに関係するため、特に重要です。

===3. Writes=== 
%8 : Tensor = aten::add_(%x0.1, %x0.1, %4) # create_model.py:20:8 
  %x0.1,

エイリアスを生成する操作は数多くあります。たとえば、次のコード 1 で AliasDb を使用すると、コード 2 に示すようなエイリアス関係がいくつか見つかります。

def forward(self, x, y): # input
x = x + 1
x.add_(x) # インプレース操作
y = y[0] # スライスまたは選択
z = [x, y] # コンテナクラス
w = torch. cat(z) はコード 1
を返します

===2. ALIAS DB===
%x.1 は次を指します: Tensor 型の WILDCARD
%8 は次を指します: %x0.1
%z.1 に含まれる: %x0.1%y0.1,
%y.1 は次を指します: 型の WILDCARD Tensor
%self は次を指します: __torch__ 型の WILDCARD.TestModel
%y0.1 は次を指します: %y.1

代码2

グラフの入力、インプレース、スライスの操作、コンテナクラスの使用によってエイリアス関係が作成され、Tensor だけがエイリアス関係を持つ可能性があるわけではないことがわかります。上の例では、z はリストです。、これは AliasDb レコードにもあります。それでは、AliasDb はどの型に注目するのでしょうか?

可変セットとワイルドカードセット

AliasDb では、変更可能な型 (可変) と変更不可能な型 (不変) の概念が導入されています。前者は Tensor や List などの内部値が変化する可能性のあるデータ型を指し、インプレース演算や追加などの操作により新しいオブジェクトを作成せずに元のオブジェクトの型を編集できます。また、int や string などの型は不変であり、AliasDb はこの部分を分析せずに単純にスキップできます。

Tuple 型は特別です。内部要素がすべて不変型 (Tuple[int] など) の場合、それも不変型になります。内部に変更可能な型 (Tuple[Tensor] など) がある場合も、不変型になります。可変型。明らかに、可変型の値が少ないほど、最適化が成功する可能性が高くなります。

可変型のオブジェクトでは、部分はを指します WILDCARD for type xxx。例として次のコードを取り上げます。

%x.1 points to: WILDCARD for type Tensor 
%y.1 points to: WILDCARD for type Tensor 
%self points to: WILDCARD for type __torch__.TestModel

この種のオブジェクトはワイルドカードセット (wildcardSet) と呼ばれます。これは、「この値の別名関係を決定できない」ことを意味します。たとえば、上記のコードでは、x と y は外部入力から取得されますが、グラフを分析するだけではストレージリソースを共有しているかどうかを判断することはできません。オブジェクトがワイルドカードセットを指すようにマークされている場合、エラーを防ぐために、それに関連する多くの最適化を回避する必要があります。

データの危険性

上記の知識があれば、より安全なパスを作成できます。パスを作成するプロセスでは、データリスクの問題 (データハザード) を解決するために AliasDb が最もよく使用されます。

例: たとえば、特定のノード A をノード B の前に挿入する場合、B が変数型パラメータ x の値を変更し、A が変更された x の値を読み取りたい場合、この挿入によりエラーが発生する可能性があります。以下のコードに示すように:

# 原图，对B的写在对A的读之前 
graph(...):  
    ... 
    B: write(x) 
    ... 
    A: read(x) 
    ... 
 
=>    # 不合法的转换！A会读取到错误的值！ 
graph(...): 
    ...  
    A: read(x) 
    B: write(x)  
    ...

このタイプの読み取り順序の変更が合法かどうかを判断するには、ノード間で読み書きされるデータ空間に重複があるかどうかを知る必要があります。getReads and 関数は AliasDb で提供されており、 Node がパラメーターとして渡され、Node が読み書きする可変変数を示す MemoryLocations オブジェクトが返されます。getWrites 上記の A と B でそれぞれgetReads and を呼び出すと、それらの間の MemoryLocation に重複がある getWrites(交差している) 場合、この交換は行われるべきではありません。次のコードに示すように:

auto loc_a = alias_db.getReads(A); 
auto loc_b = alias_db.getWrites(B); 
bool valid = !loc_a.intersects(loc_b);

このタスクをより簡単に実行できるように、AliasDb には moveAfterTopologicallyValid 関数が提供されていますmoveBeforeTopologicallyValid 。このタスクは移動前にチェックを行い、移動が合法であると判明した場合は移動されます。ここではまずワーキングセット(WorkingSet)の概念を紹介します。

WorkingSet はノードのコレクションであり、コレクション内の任意のノードは次の条件を満たします。

グラフ内でセット内の少なくとも 1 つの他のノードへの直接接続を持っています。
コレクション内の少なくとも 1 つの他のノードと MemoryLocation の読み取りと書き込みの交差が存在します (読み取りと書き込みである必要があります)。

セット外のノードと WorkingSet が上記の関係のいずれかを満たしている場合、ノードはWorkingSet に「依存」(dependOn) していると言います。

WorkingSet は合法性のチェックに役立ちます。moveAfter の例を考えてみましょう。

toMove 直前に移動したいと仮定すると movePoint 、次の 2 つの状況が考えられます。

toMovemovePoint _ の後
toMovemovePoint _ の前に

まず、を構築し WorkingSet、 toMove それを挿入し、 n が this に依存する場合は ( を含まない ) n間の WorkingSetすべてのノードをトラバースし、それを挿入する必要があります。toMovemovePointmovePointWorkingSet

ケース 1 の場合は、以下のコメントの方法に従って直接移動します。

//  `movePoint`           <dependencies>   | 
//  <dependencies>  ->    `toMove`         | `toMove` 和依赖一起移动 
//  `toMove`              `movePoint`      |

ケース 2、つまりのtoMove 前の場合はmovePoint 、 toMove 最終的に WorkingSet ここから削除します。削除する前に、 WorkingSet 合法性がチェックされます。

movePoint 依存している場合は、 WorkingSet
WorkingSet ( を含む) の toMoveノードにインプレース操作などの副作用がある場合

上記 2 つの状況のいずれかに該当する場合、移動は不正とみなされ、移動は実行されません。

合法性チェックに合格した後、以下の注釈のメソッドに従って移動されます。

//  `toMove`            `toMove`         | 
//  <dependencies>  ->  `movePoint`      | `toMove` 和依赖被分开 
//  `movePoint`         <dependencies>   |

移動に関係する WorkingSet すべてのノードにはが含まれます。このような移動は安全であり、読み取り/書き込みの競合が発生することはありません。

上記で紹介した内容は以下の関数にカプセル化されており、戻り値によって正当な移動かどうかを判断できます。

bool success = moveBeforeTopologicallyValid(A, B); 
// 如果 move 合法则进行move，返回true。否则不进行任何操作，返回false。

要約する

Jit パスの柔軟性はモデルの最適化に利便性をもたらしますが、いくつかのリスクももたらします。エイリアス分析ツール AliasDb は、これらのリスクを解決する強力なツールの 1 つです。AliasDb は MemoryDAG を使用してメモリを管理し、変更可能なデータ型と不変のデータ型を区別して、データのリスクを回避します。前の 2 つの章の知識を組み合わせると、jit モデルの生成と最適化について予備的に理解できるはずです。今後は、実践的な例から始めて、MMDeploy がこれらのツールを使用してモデルを最適化する方法を紹介していきますので、ご期待ください。

https://github.com/open-mmlab/mmdeploy github.com/open-mmlab/mmdeploy