pytorch自实现线性回归和反向传播并画图

coding=utf-8

“”"
author:lei
function:
“”"

import torch
import matplotlib.pyplot as plt

learning_rate = 0.01

1、准备数据

y = 3x + 0.8

x = torch.rand((500, 1))
y_true = x*0.3 + 0.8

2、通过模型计算y_predict

w = torch.rand((1, 1), requires_grad=True)
b = torch.empty((1, 1), requires_grad=True)

4、通过循环，反向传播，更新参数

for i in range(5000):
y_predict = torch.matmul(x, w) + b

# 3、计算loss
loss = (y_true - y_predict).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("w={}, b={}, loss={}".format(w.item(), b.item(), loss.item()))

plt.figure(figsize=(20, 8))
plt.scatter(x.numpy().reshape(-1), y_true.numpy().reshape(-1))
y_predict = torch.matmul(x, w) + b
plt.plot(x.numpy().reshape(-1), y_predict.detach().numpy().reshape(-1))
plt.show()