OpenAIは最近、勾配チェックポイントツールキットをGitHubでオープンソース化しました。このツールキットは勾配チェックポイントを設定することでメモリリソースを節約します。報告によると、一般的なフィードフォワードモデルの場合、計算時間をわずか20%増加させるだけで、GPU上で従来の10倍以上の規模のモデルを学習できるとのことです。Leifeng.comのAI研究所は、このオープンソース情報を以下のようにまとめています。

勾配チェックポイントを使用してメモリリソースを節約する

非常に深いニューラルネットワークの学習には、大量のメモリが必要です。Tim Salimans氏とYaroslav Bulatov氏が開発したgradient-checkpointingパッケージに含まれるツールは、計算時間を犠牲にすることでメモリ不足を補い、モデルをより効率的に学習できるようにします。

一般的なフィードフォワード モデルの場合、計算時間がわずか 20% 増加するだけで、GPU 上で以前より 10 倍以上大きいモデルをトレーニングすることが可能です。

ディープニューラルネットワークを学習させる際、メモリ使用量の多い部分では、バックプロパゲーションアルゴリズムを用いて損失の勾配を計算します。モデル学習中に計算グラフにチェックポイントを定義し、これらのチェックポイント間でバックプロパゲーションアルゴリズムを用いてグラフを再計算することで、メモリ使用量を削減しながら勾配値を計算できます。

n層の深層フィードフォワードニューラルネットワークを学習する場合、この手法はメモリ消費量をO(sqrt(n))まで削減できますが、追加のフォワードパス演算を実行するというコストがかかります。このライブラリは、TensorFlowでこの機能を実装できます。TensorFlowグラフエディタを使用して、バックワードパスの計算グラフを自動的に書き換えます。

図: 通常の tf.gradients 関数とこのメモリ最適化された勾配実装を使用して、異なる数のレイヤーを持つ ResNet モデルをトレーニングする場合のメモリ要件の比較。

今すぐインストールできます。

pip で tf-nightly-gpu をインストールします

pip インストール toposort networkx pytest

この手順を実行するときは、CUPTI が見つかるかどうかを確認する必要があります。

これが実行できるときです

エクスポート LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/extras/CUPTI/lib64

使い方

このライブラリは組み込み機能を提供しており、tf.gradients関数を置き換えることができます。以下のコードを入力すると、該当する関数をインポートできます。

memory_saving_gradientsからgradientsをインポートする

tf.gradients 関数を使用する場合と同様に、gradients 関数を使用してパラメーター損失の勾配を計算できます。

勾配関数には、チェックポイントという追加機能があります。

チェックポイントは、計算グラフの順方向伝播中にユーザーがグラフのどの部分をチェックしたいかを「gradients」関数に指示します。チェックポイント間のノードは、逆方向伝播中に再計算されます。

チェックポイントに一連のテンソル (gradients (ys, xs, checkpoints=[tensor1, tensor2])) を指定することも、次のキーワード (「collection」、「memory」、「speed」) を使用して設定することもできます。

tf.gradients関数をオーバーライドする

グラデーション関数を使用する別の方法は、次のように tf.gradients 関数を直接オーバーライドすることです。

TensorflowをTFとしてインポートする

メモリ節約勾配をインポートする

# モンキーパッチ tf.gradients でカスタムバージョンを指定し、チェックポイントを自動選択する

定義勾配メモリ (ys、xs、grad_ys=None、**kwargs):

memory_saving_gradients.gradients (ys, xs, grad_ys, checkpoints='memory', **kwargs) を返します。

tf.__dict__["gradients"] = 勾配メモリ

この操作の後、tf.gradients 関数を呼び出すすべてのリクエストでは、新しいメモリ節約メソッドが使用されます。

テスト

テスト フォルダーには、コードの精度とさまざまなモデルのメモリ使用量をテストするための事前に作成されたスクリプトが含まれています。

`./run_all_tests.sh` を実行してコードを変更し、テストを開始できます。

図: さまざまな数の ResNet レイヤーに対して標準勾配関数と最適化されたメモリ関数を使用した、CIFAR10 データセットでのメモリ使用量と実行時間の比較。

経由: GitHub