PyTorchで勾配降下法

PyTorchで勾配降下法をするコードを書いてみました。

最適化したい関数

def rosenbrock(x0, x1):
    y = 10 * (x1 - x0 ** 2) ** 2 + (x0 - 1) ** 2
    return y

関数を可視化する

import numpy as np

h = 0.01
x_min = -2
y_min = -3
x_max = 2
y_max = 5

X = np.arange(x_min, x_max, h)
Y = np.arange(y_min, y_max, h)

xx, yy = np.meshgrid(X, Y)

最小値はこのへんですね

matrix = rosenbrock(xx, yy)
minimum = None
min_x = None
min_y = None
for i in range(matrix.shape[0]):
    for j in range(matrix.shape[1]):
        if minimum is None or minimum > matrix[i][j]:
            minimum = matrix[i][j]
            min_y = Y[i]
            min_x = X[j]

print(min_x, min_y, minimum)

1.0000000000000027 0.9999999999999147 8.208018832734106e-26

ドットが描かれた場所が最小値になります。

import matplotlib.pyplot as plt
plt.contourf(xx, yy, np.sqrt(rosenbrock(xx, yy)), alpha=0.5)
plt.scatter(min_x, min_y, c="k")
plt.colorbar()
plt.grid()
plt.show()

勾配降下法

import numpy as np
import torch

x0 = torch.tensor(0.0, requires_grad=True)
x1 = torch.tensor(4.0, requires_grad=True)

lr = 0.001
iters = 10000

history = []
for i in range(iters):
    history.append(np.array([np.array(x0.data), np.array(x1.data)]).flatten())
    y = rosenbrock(x0, x1)
    y.backward()

    with torch.no_grad():
        x0.data -= lr * x0.grad
        x1.data -= lr * x1.grad

        x0.grad.zero_()
        x1.grad.zero_()

結果表示

import matplotlib.pyplot as plt
plt.contourf(xx, yy, np.sqrt(rosenbrock(xx, yy)), alpha=0.5)
plt.scatter([p[0] for p in history], [p[1] for p in history])
plt.scatter(min_x, min_y, c="k")
plt.colorbar()
plt.grid()
plt.show()

複数箇所からスタート

import numpy as np
import torch

x0 = torch.tensor([0.0, -0.5, -1.5, 1.5], requires_grad=True)
x1 = torch.tensor([4.0, 4.0, -1.0, -2], requires_grad=True)

lr = 0.001
iters = 10000

history = []
for i in range(iters):
    history.append([x0.detach().clone(), x1.detach().clone()])
    y = rosenbrock(x0, x1)
    #y.backward()
    s = torch.sum(y)
    s.backward()

    with torch.no_grad():
        x0.data -= lr * x0.grad
        x1.data -= lr * x1.grad

        x0.grad.zero_()
        x1.grad.zero_()

結果表示

import matplotlib.pyplot as plt
plt.contourf(xx, yy, np.sqrt(rosenbrock(xx, yy)), alpha=0.5)
for i in range(4):
    plt.scatter([p[0][i] for p in history], [p[1][i] for p in history])
plt.colorbar()
plt.scatter(min_x, min_y, c="k")
plt.grid()
plt.show()