import torch
import matplotlib.pyplot as plt

#定義訓(xùn)練集和訓(xùn)練集的答案
input_train = torch.linspace(-2, 2, steps=200)
input_train = input_train.view(200, 1)
print("input_train:", input_train)
noise = torch.normal(mean=torch.zeros(200, 1), std=0.2)
print("noise:", noise)
answer_train = torch.pow(input_train, 2) + noise
print("answer_train:", answer_train)

#定義非線性回歸模型
class NoneLinearRegression(torch.nn.Module):
? ?def __init__(self):
? ? ? ?#初始化父類
? ? ? ?super(NoneLinearRegression, self).__init__()
? ? ? ?#定義網(wǎng)絡(luò)的神經(jīng)元torch.nn.Linear(1, 10)中的1是全連接層的神經(jīng)元個(gè), 10是隱藏層神經(jīng)元的個(gè)數(shù)
? ? ? ?self.NoneLinear1 = torch.nn.Linear(1, 10)
? ? ? ?#定義神經(jīng)網(wǎng)絡(luò)的激活函數(shù)（這里使用的是雙曲正切‘tanh’激活函數(shù)）
? ? ? ?self.tanh = torch.nn.Tanh()
? ? ? ?#定義網(wǎng)貨的神經(jīng)元torch.nn.Linear(10, 1)中的10是隱藏層神經(jīng)元的個(gè)數(shù), 1是第二個(gè)全連接層的個(gè)數(shù)
? ? ? ?self.NoneLinear2 = torch.nn.Linear(10, 1)
? ?#定義前向計(jì)算
? ?def forward(self, input):
? ? ? ?#將數(shù)據(jù)傳入第一個(gè)全連接層完成計(jì)算
? ? ? ?output = self.NoneLinear1(input)
? ? ? ?#將數(shù)據(jù)傳入激活函數(shù)進(jìn)行計(jì)算
? ? ? ?output = self.tanh(output)
? ? ? ?#將數(shù)據(jù)傳入第二個(gè)全連接層進(jìn)行計(jì)算并返回計(jì)算得到的值
? ? ? ?output = self.NoneLinear2(output)
? ? ? ?return output

#定義模型
model = NoneLinearRegression()
#定義損失函數(shù)
mse_loss = torch.nn.MSELoss(reduction='mean')
#定義優(yōu)化器(注意model.parameters()一定要加(), 否則返回的不是一個(gè)可迭代的值)
optimizer = torch.optim.SGD(model.parameters(), lr=0.3)
#顯示模型中的參數(shù)
for name, parameters in model.named_parameters():
? ?print("Parameter's name is:", name)
? ?print("Parameter's value is:", parameters)

#通過訓(xùn)練集對(duì)模型進(jìn)行訓(xùn)練
for epoch in range(2000):
? ?#先進(jìn)行數(shù)據(jù)傳入進(jìn)行前向計(jì)算
? ?output_train = model(input_train)
? ?#計(jì)算損失
? ?loss = mse_loss(answer_train, output_train)
? ?#每200次查看一次損失值
? ?if (epoch + 1) % 200 == 0:
? ? ? ?print("第", epoch + 1, "次訓(xùn)練損失值為:", loss.item())
? ?#進(jìn)行梯度清零
? ?optimizer.zero_grad()
? ?#通過損失反向計(jì)算梯度
? ?loss.backward()
? ?#通過優(yōu)化器將更新的參數(shù)傳給模型
? ?optimizer.step()

#通過自己輸入一個(gè)測(cè)試集來測(cè)試模型訓(xùn)練的結(jié)果如何
input_test = input_train
output_test = model(input_test)

#通過圖像顯示出模型的訓(xùn)練效果
plt.scatter(input_train, answer_train)
#顯示模型的訓(xùn)練結(jié)果和輸入結(jié)果的擬合線
plt.plot(input_test.data.numpy(), output_test.data.numpy(), 'r-', lw=3)
plt.show()