1 import torch
 2 import matplotlib.pyplot as plt
 3 import os
 4 os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
 5 
 6 lr = 0.001
 7 see = 20000
 8 x = torch.rand([1, 50])
 9 y = 3 * x + 0.8
10 
11 w = torch.rand([1, 1], requires_grad=True, dtype=torch.float32)
12 b = torch.rand(1, requires_grad=True, dtype=torch.float32)
 
13 loss = []
14 
15 for i in range(see):
16     y_pred = torch.matmul(w, x) + b
17     cur_loss = torch.matmul(y - y_pred, (y - y_pred).T)18     loss.append(cur_loss.item())
19 
20     if i != 0:  # 將梯度清零，初始時引數的梯度為None所以先計算一次後才有梯度
21         w.grad.data.zero_()
22         b.grad.data.zero_()
23 
24     cur_loss.backward()
 
25     w.data = w.data - lr * w.grad
26     b.data = b.data - lr * b.grad
27 
28     if i % 200 == 0:
29         print("w, b, loss", w.item(), b.item(), cur_loss.item())
30 
31 plt.scatter(x.numpy()[0], y.numpy()[0])
32 y_predict = torch.matmul(w, x) + b
33 plt.plot(x.numpy()[0], y_predict.detach().numpy()[0])
 
34 plt.show()