import matplotlib.pyplot as plt # 导入画图的程序包
import torch
import numpy as np
x = torch.linspace(0,100,100).type(torch.FloatTensor)
rand = torch.randn(100)* 10
y = x + rand
#划分训练集与测试集。90个训练 10个测试。
x_train = x[ : -10]
x_test = x[-10 : ]
y_train = y[ : -10]
y_test = y[-10 : ]

# plt.figure(figsize=(10,8)) #设定绘制窗口大小10*8
# plt.xlabel('X') #添加X轴标注
# plt.ylabel('Y') #添加Y轴的标注
# plt.plot(x_train.numpy(),y_trian.numpy(),'o')
# plt.show() #画出图形

#定义两个自动微分变量a和b
a = torch.rand(1, requires_grad=True)
b = torch.rand(1, requires_grad=True)
#设置学习率
learning_rate = 0.0001
"""对a,b进行迭代计算。"""
for i in range(1000):
#计算当前a,b条件下的模型预测值
    predictons = a.expand_as(x_train) * x_train + b.expand_as(x_train)
    #将所有训练数据带入模型ax+b，计算每个预测值。这里的x_train和predictions都是（90，1）的张量
    loss = torch.mean((predictons - y_train) ** 2)
    print('loss:', loss)
    loss.backward()

    a.data = a.data - learning_rate * a.grad
    b.data = b.data - learning_rate * b.grad
    a.grad.data.zero_()  # 清空a的梯度数值
    b.grad.data.zero_()

x_data = x_train.data.numpy() # 获得x包裹的数据
x_pred = x_test.data.numpy() # 获得包裹的测试数据的自变量
plt.figure(figsize = (10, 7)) # 设定绘图窗口大小
plt.plot(x_data, y_train.data.numpy(), 'o') # 绘制训练数据
plt.plot(x_pred, y_test.data.numpy(), 's') # 绘制测试数据
x_data = np.r_[x_data, x_test.data.numpy()]
xplot = plt.plot(x_data, a.data.numpy() * x_data + b.data.numpy())  # 绘制拟合数据
yplot = plt.plot(x_pred, a.data.numpy() * x_pred + b.data.numpy(), 'o') # 绘制预测数据
plt.xlabel('X') # 更改横坐标轴标注
plt.ylabel('Y') # 更改纵坐标轴标注
str1 = str(a.data.numpy()[0]) + 'x +' + str(b.data.numpy()[0]) # 图例信息
plt.legend([xplot, yplot],['Data', str1]) # 绘制图例
plt.show()