PyTorch autograd で Newton法実装

PyTorchは主にニューラルネットワークの学習に使用されるライブラリですが, autogradはより幅広い使い方ができます. 今回はその機能を用いてNewton法を実装しました.

autogradとは

autograd機能は以下のようなものです.

PyTorchのニューラルネットワークはautogradパッケージが中心になっています．autogradは自動微分機能を提供します．つまり，このパッケージを使うと勝手に微分の計算を行ってくれると言うことです．(引用: [1] 【PyTorch入門】第2回 autograd:自動微分).

次の記事も参考になります ([2] PyTorch公式チュートリアル Deep Learning with PyTorch #2 Autograd).

Newton法とは

f(x) = 0

の解を求める, 反復球根アルゴリズムの一つです ([3] ニュートン法).

更新式は

x_{n+1} = x_{n} - \frac{f(x_n)}{f'(x_n)}

です, $f'(x_n)$が陽に求められない場合は, 近似値を用いることもあります.

実装(本質的な部分)

計算に必要な本質的な部分を紹介します. 以下では例として $f(x) = x^2$を扱います.

主に使用するのはVariableクラスです. 上記の更新式$x_{n}$に相当する変数xを定義します.

import torch
from torch.autograd import Variable #自動微分機能のimport

x = Variable(tensor, requires_grad=True)

第1引数tensorにはtensorクラスの変数を入力します. 例えば, 初期値x=5を入力する場合は以下のように書きます.

ini_x = 5*torch.ones(1,1) # ones(1,1)で値が1の1×1のtensor(スカラー値)が生成される
x = Variable(ini_x, requires_grad=True)
>> tensor([[5.]])

オプションrequires_gradで, この変数を微分の対象にするかを指定します.

次に

f = x*x
f.backward() # 勾配の計算

にて$f(x)=x^2$の微分が行われ, x.gradにxの値での勾配が保存されます.

x.grad # 勾配の表示
>> tensor([[10.]])

f'(x) = 2*x なので f'(5) = 10 が再現されていることが分かります.

また, これは注意事項ですが
backwardにより計算される勾配は, 実行するたびに累積されるため,
反復計算を行う場合は, 勾配を0に手動で設定する必要があります ([4] Why do we need to set the gradients manually to zero in pytorch?).

x.grad.zero_() # 勾配を0に設定
>> tensor([[0.]])

いくつかの知識

Variableクラスについて,

• VariableはTensorクラスをラップしており, 生のtensorはx.dataでアクセス可能です([2]).
• 計算された勾配x.gradもtensor型で保存されます.
• またx.dataに保存された値はx.item()により取り出すことができます.

実装

x = \sqrt{2}

となる, $x$を求めます. これは

f(x) = x^2 -2

として, $f(x) = 0$ となるxを求めれば良いですね.

import torch
from torch.autograd import Variable #自動微分機能のimport

# f(x)の定義
def f(x):
    return x*x-2

# 初期点の設定
ini_x_value = 5

ini_x_tensor = ini_x_value*torch.ones(1, 1, dtype = torch.float64)
x = Variable(ini_x_tensor, requires_grad=True)
print(f'roop {0:<4d} x = {x.item()}')

roop = 0
while roop < 10:
    roop += 1
    # 勾配の計算
    f(x).backward()
    # xの更新
    x.data -= (f(x)/x.grad).data
    # 勾配を0に設定
    x.grad.zero_()
    print(f'roop {roop:<4d} x = {x.item()}')

いくつかの注意

• 反復計算中, x自体ではなく, x.dataのみを更新します. x -= f(x)/x.grad のように x自体を更新すると, grad属性がNoneになり次のループでの計算時にエラーが生じるからです.
• f(x).backward() は (x*x-2).backward() と同等です.
• 前述したように, 次の勾配を計算する前にx.grad.zero_())で勾配を0に手動で設定する必要があります.(重要なのでもう一度書きました)

出力結果

roop 0    x = 5.0
roop 1    x = 2.7
roop 2    x = 1.7203703703703703
roop 3    x = 1.44145536817765
roop 4    x = 1.414470981367771
roop 5    x = 1.4142135857968836
roop 6    x = 1.4142135623730951
roop 7    x = 1.414213562373095
roop 8    x = 1.4142135623730951
roop 9    x = 1.414213562373095
roop 10   x = 1.4142135623730951

と, $\sqrt{2} \approx 141421356237309504880 \ldots $ Decimal expansion of square root of 2 に近い値を求めることができていますね.

*** 2020/02/04 (編集)
@nabeyyy さまからのご指摘により
ini_x_tensor = ini_x_value*torch.ones(1, 1) を
ini_x_tensor = ini_x_value*torch.ones(1, 1, dtype = torch.float64) に変更しました.

まとめ

newton法を通してautogradの使い方を紹介しました.

参考

[1] @manaco
, "【PyTorch入門】第2回 autograd:自動微分", https://qiita.com/manaco/items/f4be3fb0d996a6a3eae3)
[2] @asai0304,
PyTorch公式チュートリアル Deep Learning with PyTorch #2 Autograd, https://qiita.com/asai0304/items/e6f413aa20b927e4ebf6
[3] "ニュートン法", https://ja.wikipedia.org/wiki/ニュートン法
[4] "Why do we need to set the gradients manually to zero in pytorch?", https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903