搭建一元一次线性回归方程（自定义类）

创建了数据集x和y，通过自定义模型，创建忧化函数，损失函数，通过大量训练能过得出预测的y值,当x=4时>>>7.9998 和真实的y=8 十分接近。

import torch
from torch import nn
from matplotlib import pyplot as plt
import torch.optim as optim


#创建数据集
x=torch.tensor([[1.0],[2.0],[3.0]])
y=torch.tensor([[2.0],[4.0],[6.0]])


#自定义线性回归模型
#让LinearRegression继承Moudule的内容，Moudule是创建线性回归方程的一个类。
class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression,self).__init__()
        self.fc=nn.Linear(1,1)  #(1，1)表示定义一个一元一次线性回归方程
    def forward(self,x):
        y_pred=self.fc(x)
        return y_pred
#创建完回归方程类后，再把它创建出来。
model=LinearRegression()
#创建损失函数和优化函数
criterion=nn.MSELoss(reduction='sum')  #创建损失函数
optimizer=optim.SGD(model.parameters(),lr=0.01)#创建优化函数


#模型训练
e_list=[]
l_list=[]
w_list=[]
b_list=[]
for i in range(1000):
    y_pred=model(x)
    loss=criterion(y_pred,y)
    e_list.append(i)
    l_list.append(loss)
    w_list.append(model.fc.weight.item())  #返回系数的数值
    b_list.append(model.fc.bias.item())
    optimizer.zero_grad() #梯度归零
    loss.backward()    #回车一下，把loss的值返回一下，应该相当于return
    optimizer.step()   #更新

#测试模型
x_test=torch.tensor([[4.0]])
y_test=model(x_test)
print(y_test)

原文地址：http://www.cnblogs.com/JK8395/p/16881178.html