Pytorchで単回帰やってみる

先日の記事を元に、単回帰(simple regression)をpytorchで書いてみます。

当然ですが、型とか結構大事。。

まずは設定

import torch
from numpy.random import *
import numpy as np

import matplotlib.pyplot as plt

次にデータを作ってみます

話をわかりやすくするために、

$$y = 2x + 3$$

に少し乱数を振っておきましょうか。

X = np.array([[x * 10] for x in rand(20)], dtype='float32')
y = np.array([2 * x[0] + 3 + 0.3 * randn(1, 1)[0] for x in X], dtype='float32')

plt.scatter(X, y)

これらデータを使って、

$$y = Wx + b$$

と置いて、 $W, b$ を推定します（点推定）。勿論、最終的に $W=2, b=3$ となればOKですね。
(無理やり？)NNモデルと考えるので、入力層、出力層を定義しなければいけませんが、通常はベクトルですがココでは実質スカラーですね。ということで、

input_size, output_size = 1,1

で、推定する傾き$W$をpytorchから見えるように定義します。

w = torch.rand(input_size, hidden_size, require_grad=True)
b = torch.rand(1, requires_grad=True)

あとは、学習率(パラメータ$W, b$の更新幅)を設定しますね。

learning_rate = 1e-6

pytorch から見えるように元データを変換しておきます。

X_train = torch.from_numpy(X)
Y_train = torch.from_numpy(y)

いざ学習

epochs = 10000

for iter in range(epochs):
    ## 予測させます(feed forward)
    y_pred = X_train.mm(w) + b
    ## 誤差を計算
    loss = (y_pred - Y_train).pow(2).sum()
    if iter % 10000 == 0:
        print(str(iter) + ','+ str(loss.item()))
    ## 逆誤差伝搬で勾配を計算
    loss.backward()
    with torch.no_grad():
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad
        w.grad.zero_()
        b.grad.zero_()

こうやって計算させると、

w, b
#(tensor([[1.9053]], requires_grad=True), tensor([3.4280], requires_grad=True))

なので、元の値 $W=2, b=3$ にそれなりに近い値が求められました。

## 図示用に変形
x_train_tensor = torch.from_numpy(X)
predicted_in_tensor = x_train_tensor.mm(w) + b
predicted = predicted_in_tensor.detach().numpy()

## 図示
plt.scatter(X, y, label='original data', s=20, c='g')
plt.plot(X, predicted, label='fitted line')
plt.legend()