本节实现一个训练线性回归参数的例子，线性回归的损失函数为：$loss=\frac{1}{2} \sum_{i=1}^{N}(y_{i}-(wx_{i}+b))^2$，然后利用随机梯度下降法更新参数 $w$ 和 $b$ 来最小化损失函数，最终学得 $w$ 和 $b$ 的值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import torch as t
from matplotlib import pyplot as plt

device = t.device('cuda:0')

# 设置随机数种子，保证在不同电脑上运行时下面的输出一致
t.manual_seed(1000)


def get_fake_data(batch_size=8):
    x = t.rand(batch_size, 1, device=device) * 5
    y = x * 2 + 3 + t.randn(batch_size, 1, device=device)
    return x, y


# 随机初始化参数
w = t.rand(1, 1).to(device)
b = t.zeros(1, 1).to(device)

lr = 0.02  # 设置学习率

for ii in range(1000):
    x, y = get_fake_data(batch_size=4)

    # forward: 计算loss
    y_pred = x.mm(w) + b.expand_as(y)
    loss = 0.5 * (y_pred - y) ** 2
    loss = loss.mean()

    # backward: 手动计算梯度
    dloss = 1
    dy_pred = dloss * (y_pred - y)

    dw = x.t().mm(dy_pred)
    db = dy_pred.sum()

    # 更新参数
    w.sub_(lr * dw)
    b.sub_(lr * db)

    if ii % 50 == 0:
        # 画图
        x = t.arange(0, 6).view(-1, 1)
        x = x.float()
        y = x.mm(w.cpu()) + b.cpu().expand_as(x)
        plt.plot(x.cpu().numpy(), y.cpu().numpy())  # predicted

        x2, y2 = get_fake_data(batch_size=32)
        plt.scatter(x2.cpu().numpy(), y2.cpu().numpy())  # true data

        plt.xlim(0, 5)
        plt.ylim(0, 13)
        plt.show()
        plt.pause(0.5)

print('w: ', w.item(), 'b: ', b.item())

笔记来源：《pytorch-book》