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()
