pytorchの線形回帰を練習

ここ数日、PyTorch for Deep Learning - Full Course / Tutorialという動画を見ています。
この動画を見終えると、GANで画像を生成する技術が身につくそうです。

PyTorch for Deep Learning - Full Course / Tutorial: https://youtu.be/GIsg-ZUy0MY
GANでの画像生成のすごいやつ：
　GANを使って簡単に架空アイドル画像を自動生成（Progressive Growing of GANs）https://www.mahirokazuko.com/entry/2018/12/15/201501

この9時間のコースの第2章、"Linear Regression"の内容が頭にまだ入りきっていないので、qiitaの記事を書くことで定着を図ろうとしています。

目標

以下で行っていることを、pytorchを使って表現できるようになることです。
できれば、サンプル数を900000などに設定しなくても、うまくいってほしいです。

import random

# a * x + c の a, cを予測する。

def two_x_plus_three(x):
    return 2 * x + 3
#入力をランダムに作る。
inputs = [random.randint(0, 100) for _ in range(900000)]

#入力の個数にあわせて、targetを作る。
targets = [two_x_plus_three(x) for x in inputs]

#学習される前のa, cの予測値をランダムに作る。
a_assumed = float(random.randint(0, 100))
c_assumed = float(random.randint(0, 100))

#学習率
learning_rate = 0.01

for i, x in enumerate(inputs):
    #e: 微分に使う微小な値
    e = 0.001

    #pred: 予測値
    pred = a_assumed * x + c_assumed
    #予測値と目標値の差を表す関数diffを、a_assumed, c_assumedについて微分する。
    pred_ae = (a_assumed - e) * x + c_assumed
    pred_ce = a_assumed * x + (c_assumed - e)

    diff = (pred - targets[i]) ** 2
    diff_da = diff - (pred_ae - targets[i]) ** 2
    diff_dc = diff - (pred_ce - targets[i]) ** 2

    #得られた傾きと逆方向に、a_assumed, c_assumedを動かす。
    a_assumed -= diff_da * learning_rate
    c_assumed -= diff_dc * learning_rate

print(a_assumed) #1.99など
print(c_assumed) #3.37など

Gradient Descentを用いた学習の6つのステップ

学習は、以下の6つのステップをたどります。

0 データセットを用意する
1 予測を生成する
2 損失を計算する
3 重さとバイアスに関連する勾配を計算する
4 学習率の分だけ、勾配を重さとバイアスに反映する
5 勾配をゼロに戻す

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import random

# 0 データセットを用意する

def two_x_plus_three(x):
    return 2 * x + 3

inputs = np.array(\
         [[random.randint(0, 100)] for _ in range(300)],\
         dtype = 'float32')

targets = np.array(\
          [two_x_plus_three(x) for x in inputs],\
          dtype = 'float32')

inputs  = torch.from_numpy(inputs)
targets = torch.from_numpy(targets)

train_ds = TensorDataset(inputs, targets)
batch_size = 5
#DataLoaderを使うと、batch_sizeの分だけ毎回学習を行える。
train_dl = DataLoader(train_ds, batch_size, shuffle = True)

# 1 予測を生成する　に使う
model = nn.Linear(1,1)

# 2 損失を計算する, 3 勾配を計算する　に使う
loss_fn = F.mse_loss
loss = loss_fn(model(inputs), targets)

# 4 学習率の分だけ、勾配を重さとバイアスに反映する　に使う
opt = torch.optim.SGD(model.parameters(), lr = 0.0001)

def fit(num_epochs, model, loss_fn, opt, train_dl):
    for epoch in range(num_epochs):

        for x_batch, y_batch in train_dl:
            # 1 予測を生成する。
            pred = model(x_batch)

            # 2 損失を計算する
            loss = loss_fn(pred, y_batch)

            # 3 勾配を計算する
            loss.backward()

            # 4 学習率の分だけ、勾配を重さとバイアスに反映する
            opt.step()

            # 5 勾配をゼロに戻す
            opt.zero_grad()

fit(2000, model, loss_fn, opt, train_dl)

print(list(model.parameters()))

結果：

[Parameter containing:
tensor([[2.0002]], requires_grad=True), Parameter containing:
tensor([2.9884], requires_grad=True)]

y =2x + 3なので、なかなかいいんじゃないでしょうか。