一、前言
简单的使用pytorch拟合一条直线
二、实现
import torch
import matplotlib.pyplot as plt
# 学习率
learning_rate = 0.01
# 数据集(500条训练, 50条验证)
train_x = torch.rand([500, 1])
train_y = train_x*2 + 1
val_x = torch.rand([50, 1])
val_y = val_x*2 + 1
print(f"train_x: {train_x}")
print(f"train_y: {train_y}")
print(f"val_x: {val_x}")
print(f"val_y: {val_y}")
# 随机初始化网络权重(w, b)
w = torch.rand([1, 1], requires_grad=True, dtype=torch.float32) # requires_grad=True 表示该参数需要计算梯度
b = torch.tensor(0, requires_grad=True, dtype=torch.float32) # requires_grad=True 表示该参数需要计算梯度
print(f"w: {w}")
print(f"b: {b}")
# 初始化画布(绘制训练精度)
fig = plt.figure(figsize=(5, 4))
# 训练
for i in range(2000):
# 前向传播
predict_y = torch.matmul(train_x, w) + b
# 计算损失(均方误差)
loss = (train_y - predict_y).pow(2).mean()
# 梯度清零
if w.grad is not None:
w.grad.data.zero_()
if b.grad is not None:
b.grad.data.zero_()
# 反向传播,计算梯度
loss.backward()
# 更新梯度
w.data = w.data - learning_rate*w.grad
b.data = b.data - learning_rate*b.grad
# 打印损失值
print(f"epochs:{i + 1}, loss: {loss}")
# 绘制训练精度
if i % 50 == 0 and i != 0: # 每隔50步校验精度并绘制图像
with torch.no_grad(): # 验证时不记录梯度
# 绘制验证集真实值(直线)
plt.plot(val_x.numpy().reshape(-1), val_y.numpy().reshape(-1))
# 送入验证集,前向传播计算验证集预测结果
test_y = torch.matmul(val_x, w) + b
# 绘制验证集预测值(散点)
plt.scatter(val_x.numpy().reshape(-1), test_y.detach().numpy().reshape(-1), c='g')
plt.draw() # 绘制
plt.pause(0.1) # 暂停
fig.clf() # 清除当前帧