AdaGradを使った損失関数の最小化問題をPythonで実装

import numpy as np
import matplotlib.pyplot as plt
import math

class AdaGrad:
  def __init__(self,lr):
    self.lr = lr
    self.h = None
  
# パラメータの更新
  def update(self,grad,x_initial):
    if self.h is None:
      self.h = 0

    self.h = self.h + grad**2
    x_next = x_initial - self.lr * grad / (math.sqrt(self.h)+1e-7)
    return x_next

# 損失関数の設定
def loss_function(x):
  L = 5*x**2 + 3*x
  return L

# 勾配の計算
def calc_gradient(x):
  grad = 10*x + 3
  return grad

# xの初期値
x_initial = -5

# 損失値
L = loss_function(x_initial)

# 損失関数をプロットする
x_line = np.linspace(-10, 20) 
L_line = loss_function(x_line)
plt.figure(figsize=(6,6))
plt.plot(x_line, L_line)
plt.xlabel('parameter x')
plt.ylabel('loss function L')
# 現在の x と損失値をプロット
plt.scatter(x_initial, L, color='r')
plt.xlim([-10, 10])
plt.ylim([-5, 300])
plt.show()

from matplotlib import animation, rc
from IPython.display import HTML

# 学習率(lr = learning rate)
lr = 2

# AddGradオブジェクトの宣言
adagrad = AdaGrad(lr)

# 更新回数
num_iterations = 50

# 初期値をコピー
x = x_initial

# 描画
fig = plt.figure(figsize=(6,6))
plt.plot(x_line, L_line)
images = []

for n in range(num_iterations):
    # 損失値を計算
    L = loss_function(x)
    # 勾配を計算
    grad = calc_gradient(x)

    print("%d-th iteration, x=%.3f, loss: %.3f, grad: %.3f" % (n, x, L, grad))

    # 更新前の状態を描画
    tangent = grad * (x_line - x) + L
    img = plt.plot(x_line, tangent, color='r')
    img.append(plt.scatter(x, L, color='r'))
    img.append(plt.text(-8, 300, 'iteration: '+str(n), size='x-large'))
    images.append(img)
    
    # 更新
    x = adagrad.update(grad,x)

plt.xlim([-10, 10])
plt.ylim([-5, 400])
plt.xlabel('parameter x')
plt.ylabel('loss function L')

# アニメーション作成
anim = animation.ArtistAnimation(fig, images, interval=100)

# Google Colaboratoryの場合必要
rc('animation', html='jshtml')
plt.close()
display(anim)