使用torch 进行梯度下降

是针对的课后作业 来实现的 老师的课听了好几遍才大致理解整个的流程

import torch
import matplotlib.pyplot as plt

x_data = [1.0,2.0,3.0]
y_data = [2.0,4.0,6.0]

w1 = torch.Tensor([1.0])
w2 = torch.Tensor([1.0])
b = torch.Tensor([1.0])

#此代码为了 对相应参数进行梯度下降并保留到相关参数内
w1.requires_grad = True
w2.requires_grad = True
b.requires_grad = True

def forward(x):
    return w1*x**2 + w2*x + b

def loss(x,y):
    y_pred = forward(x)
    loss = (y_pred-y)**2
    return loss

print("before training", 4, forward(4))

epoch_list = []
loss_list = []

for epoch in range(10001):
    for x,y in zip(x_data,y_data):
        l = loss(x, y)
        l.backward()   # 计算出所有相关梯度分别存于 w1.grad,w2.grad,b.grad中  使用的是SGD
        w1.data = w1.data - 0.01*w1.grad.data
        w2.data = w2.data - 0.01*w2.grad.data
        b.data = b.data - 0.01*b.grad.data
        
        #torch 的每次运算都在构建运算图，在运算完一次循环后忘记加下面的梯度清除了这个了 导致数值非常大 出现了nan
        w1.grad.data.zero_()
        w2.grad.data.zero_()
        b.grad.data.zero_()
        
    if epoch %1000 == 0:  #每1000步打印一次相关数据
        print("epoch", epoch, "loss:",  w1.data, w2.data, b.data, l.item())
    epoch_list.append(epoch)
    loss_list.append(l)
        
print("after training", 4, forward(4))

plt.plot(epoch_list,loss_list)
plt.ylabel("Loss")
plt.xlabel("epoch")
plt.show()