PyTorchが計算グラフを作成する仕組み

小さな小さなPyTorchクローンを作ろうと思ったので、PyTorchの計算グラフの仕組みを調べました。

公式ブログに分かりやすいGIFアニメがあったので、メモに残します。

How Computational Graphs are Constructed in PyTorch

１．requires_gradなテンソルを関数に入力

２．grad_fn（勾配関数）ノードが作成される

３．「collect_next_edges」が入力関数からgrad_fnにリンクするエッジを作成

６．同様にして関数Bを作成

７．３つのテンソルを関数Cへ入力

９．「collect_next_edges」が入力の関数を調べて、grad_fnへのエッジを作成する

１０．終了。