ゼロから作るDeep Learning ❸ ―フレームワーク編

ステップ18メモリ使用量を減らすモード

前ステップでは、Pythonのメモリ管理について学びました。ここでは、DeZeroのメモリ使用を改善できる仕組みを2つ取り入れます。1つ目の改善点は、逆伝播で消費するメモリ使用量についてです。そこでは、不要な微分をすぐに消去する仕組みを用意します。そして2つ目の改善点は、「逆伝播が必要ない場合のモード」を用意することです。そのモードでは不要な計算は省くようにします。

18.1　不要な微分は保持しない

1つ目の改善点は、DeZeroの逆伝播についてです。現在、私たちのDeZeroでは、すべての変数が微分を保持します。たとえば、次の例を見てみましょう。

x0 = Variable(np.array(1.0))
x1 = Variable(np.array(1.0))
t = add(x0, x1)
y = add(x0, t)
y.backward()

print(y.grad, t.grad)
print(x0.grad, x1.grad)

実行結果

1.0 1.0
2.0 1.0

ここで、ユーザが用意した変数はx0とx1です。その他の変数tとyは計算によって生成されます。このときy.backward()によって微分を求めると、すべての変数は微分を保持します。しかし、多くの場合――特に機械学習では――、逆伝播によって求めたい微分は、末端の変数（x0、x1）だけの場合がほとんどです。上の例では、yやtなどの中間の変数については、その微分を必要とするケースは多くありません。そこで、それらの中間の変数については微分を消去するモードを追加します。そのために、現状のVariableクラスのbackwardメソッドに対して、次の網掛け部分のコードを追加します。 ...