TensorFlowの自動微分を使って勾配降下法を試してみる

TensorFlowの自動微分を使ってみるために前回の記事(勾配降下法の自作アルゴリズム)に書いたアルゴリズムをTensorFlowで書きました。

 \displaystyle  z = (x^2 + y^2 - 1)^2 + x の最小値となる  \displaystyle  x  \displaystyle  y を求めることにします。

この関数はxyの座標でいうと、原点を中心に半径1の円から少し左にずれた環状が谷になっていて、その中でも左のほうが低くなっております。

f:id:suzuki-navi:20200816235427p:plain

f:id:suzuki-navi:20200816235441p:plain

以下は、TensorFlowの自動微分を使って勾配降下法で最小となるxyを求めたコードです。

Google Colaboratoryで実行しました。

普通の勾配降下法

import numpy as np
import matplotlib.pyplot as plt
import math
import tensorflow as tf
# 適当な初期値
x = tf.Variable(1.5)
y = tf.Variable(1.1)

# 目的となる損失関数
def loss():
  x2 = x * x
  y2 = y * y
  r2 = (x2 + y2 - 1.0)
  return r2 * r2 + x

# グラフにするための配列
counterHistory = []
xHistory = []
yHistory = []
lossHistory = []

varlist = [x, y]

def loop():
  for loopCount in range(500):
    with tf.GradientTape() as tape:
      l = loss()
    counterHistory.append(loopCount)
    xHistory.append(float(x))
    yHistory.append(float(y))
    lossHistory.append(float(l))
    d = tape.gradient(l, varlist)
    for i in range(len(varlist)):
      varlist[i].assign(varlist[i] - 0.01 * d[i])

loop()

print(x) # -1.1053842
print(y) # 0.059747234

# グラフ化1つ目
plt.plot(counterHistory, xHistory)
plt.plot(counterHistory, yHistory)
plt.plot(counterHistory, lossHistory)
plt.show()

# グラフ化2つ目
plt.xlim(-2, +2)
plt.ylim(-1.5, +1.5)
plt.plot(xHistory, yHistory)
th = np.linspace(-math.pi, math.pi, 100)
plt.plot(np.cos(th), np.sin(th))
plt.show()

1つ目のグラフは、ステップを踏むごとに最適化される様子です。

  • 緑: 損失関数
  • 青: x
  • 橙: y

f:id:suzuki-navi:20200816235825p:plain

2つ目のグラフは、xとyの推移をxy座標にプロットしたものです。

  • 青: xy
  • 橙: 半径1の円

f:id:suzuki-navi:20200816235804p:plain

0.01という学習率を使って、500回ループを回して、まだわずかに最小値には達していないように見えます。学習率を少し大きくしたらもっと早く収束しましたが、大きすぎると収束しませんでした。

学習率0.2: 大きく振幅して収束せず 学習率0.15: 少し振幅するが、100回ループで最小値に十分近づく 学習率0.08: 少し振幅するが、100回ループで最小値にある程度近づく

前回記事のアルゴリズム

# 適当な初期値
x = tf.Variable(1.5)
y = tf.Variable(1.1)

# 目的となる損失関数
def loss():
  x2 = x * x
  y2 = y * y
  r2 = (x2 + y2 - 1.0)
  return r2 * r2 + x

# グラフにするための配列
counterHistory = []
xHistory = []
yHistory = []
lossHistory = []

varlist = [x, y]

def loop():
  k = math.log(3.0)
  prev_grad = []
  prev_diff = []
  for loopCount in range(40):
    with tf.GradientTape() as tape:
      l = loss()
    counterHistory.append(loopCount)
    xHistory.append(float(x))
    yHistory.append(float(y))
    lossHistory.append(float(l))
    d = tape.gradient(l, varlist)
    curr_grad = []
    curr_diff = []
    if loopCount == 0:
      for i in range(len(varlist)):
        curr_grad.append(d[i].numpy())
        curr_diff.append(curr_grad[i] * -0.01)
        varlist[i].assign(varlist[i] + curr_diff[i])
    else:
      for i in range(len(varlist)):
        curr_grad.append(d[i].numpy())
        z = curr_grad[i] / prev_grad[i]
        alpha = -0.5 + 2.0 / (1.0 + np.exp(k * (1.0 - 3.0 * z)))
        curr_diff.append(alpha * prev_diff[i])
        varlist[i].assign(varlist[i] + curr_diff[i])
    prev_grad = curr_grad
    prev_diff = curr_diff

loop()

print(x) # -1.1071593
print(y) # -1.0205355e-07

# グラフ化1つ目
plt.plot(counterHistory, xHistory)
plt.plot(counterHistory, yHistory)
plt.plot(counterHistory, lossHistory)
plt.show()

# グラフ化2つ目
plt.xlim(-2, +2)
plt.ylim(-1.5, +1.5)
plt.plot(xHistory, yHistory)
th = np.linspace(-math.pi, math.pi, 100)
plt.plot(np.cos(th), np.sin(th))
plt.show()
  • 緑: 損失関数
  • 青: x
  • 橙: y

f:id:suzuki-navi:20200816235614p:plain

  • 青: xy
  • 橙: 半径1の円

f:id:suzuki-navi:20200816235630p:plain

初回だけ0.01という学習率を使って計算し、2回目以降を前回記事のアルゴリズムで計算しました。40回のループを回して、十分に収束しました。学習率はいくつにしてもそんなに変わらないです。

2020/09/22 追記

このアルゴリズムをTensorFlowのOptimizerで実装しました。