import numpy as np
import matplotlib.pyplot as plt
from torch import nn, optim
from torch.autograd import Variable
import torch


x_data = np.random.rand(100)
# 加入高斯噪声, 且和x_data形状一致
noise = np.random.normal(0, 0.01, x_data.shape)

y_data = x_data*0.1 + 0.2 + noise

# plt.scatter(x_data, y_data)
# plt.show()

# 将数据从一维数据传换成二维的数据
""" -1 表示自动匹配行 """
x_data = x_data.reshape(-1, 1)
y_data = y_data.reshape(-1, 1)


# 把 numpy 转化成 tensor 数据
x_data = torch.FloatTensor(x_data)
y_data = torch.FloatTensor(y_data)

inputs = Variable(x_data)
target = Variable(y_data)


# 构建神经网络模型
# 一般把网络中具有可学习参数的层放在__init__()中
class LinearRegression(nn.Module):
    # 定义网络结果
    def __init__(self):
        # 初始化 nn.Module
        super(LinearRegression, self).__init__()
        # 定义网络结构, 神经元个数
        self.fc = nn.Linear(1, 1)

    # 定义网络计算
    def forward(self, x):
        out = self.fc(x)
        return out

# 定义模型
model = LinearRegression()

# 定义代价函数
mse_loss = nn.MSELoss()

# 定义优化器, lr表示学习率
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 分别为参数的名字和参数的数值
for name, parameters in model.named_parameters():
    print('name:{}, param:{}'.format(name, parameters))

# 模型训练
for i in range(1000):
    out = model(inputs)
    # 计算loss
    loss = mse_loss(out, target)

    # 梯度清 0
    optimizer.zero_grad()

    # 计算梯度, 不清零会导致梯度累加
    loss.backward()

    # 修改权值
    optimizer.step()

    if i % 200 == 0:
        print(i, loss.item())

# 查看回归线的准确度
y_pred = model(inputs)
plt.scatter(x_data, y_data)
plt.plot(x_data, y_pred.data.numpy(), 'r-', lw=3)
plt.show()