PyTorch示例——利用PyTorch计算梯度

版本信息

• PyTorch: 1.12.1
• Python: 3.7.13

导包

import torch
import matplotlib.pyplot as plt

原始数据

# 原始数据
x_data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
y_data = [0.4, 1.4, 2.8, 3.5, 4.8, 5.5]
plt.scatter(x_data, y_data)
plt.show()

定义函数

def forward(x):
    # 前向传播
    return x * w

def loss(out, y):
    # 计算损失
    return (out - y) ** 2

计算梯度

• 代码

# 权重w，需要计算梯度
w = torch.Tensor([10.0])
w.requires_grad = True
# 学习率lr，不需要计算梯度
lr = torch.Tensor([0.005])
lr.requires_grad = False

loss_list = []
epochs = range(10)
for epoch in epochs:
    for x, y in zip(x_data, y_data):
        out = forward(x)
        l = loss(out, y)
        l.backward() # 反向传播计算每个点的梯度
        w.data = w.data - lr * w.grad.data # 学习，梯度更新
        w.grad.data.zero_() # 重置本次累积的梯度
    print(f"epoch = {epoch}, loss_val = {l.item()}")
    loss_list.append(l.item())

• 输出结果

epoch = 0, loss_val = 881.4310302734375
epoch = 1, loss_val = 107.05437469482422
epoch = 2, loss_val = 12.973189353942871
epoch = 3, loss_val = 1.5620064735412598
epoch = 4, loss_val = 0.18457315862178802
epoch = 5, loss_val = 0.020624106749892235
epoch = 6, loss_val = 0.0019251183839514852
epoch = 7, loss_val = 8.275721484096721e-05
epoch = 8, loss_val = 9.185609087580815e-06
epoch = 9, loss_val = 5.270535984891467e-05

绘制曲线：epoch与loss

plt.plot(epochs, loss_list)
plt.xlabel("epoch")
plt.ylabel("loss_val")
plt.show()