线性回归的pytorch代码

使用pytorch实现的线性回归， 闲言少叙，直接上代码，客官请看：

import torch
import torch.nn as  nn
import numpy as np
import matplotlib.pyplot as plt

#设置相关参数
input_size=1
output_size=1
num_epochs=60
learning_rate=0.001

#导入训练数据集
x_train=np.array([[3.3],[4.4],[5.5],[6.71],[6.93],[4.168],
                  [9.779],[6.182],[7.59],[2.167],[7.042],
                  [10.791],[5.313],[7.997],[3.1]],dtype=np.float32)
y_train=np.array([[1.7],[2.76],[2.09],[3.19],[1.694],[1.573],
                  [3.366],[2.596],[2.53],[1.221],[2.827],
                  [3.465],[1.65],[2.904],[1.3]],dtype=np.float32)

#设置模型、损失函数、优化函数等
model=nn.Linear(input_size,output_size)
criterion=nn.MSELoss()
optimizer=torch.optim.SGD(model.parameters(),lr=learning_rate)

#开始迭代训练
for epoch in  range(num_epochs):
    inputs=torch.from_numpy(x_train)
    targets=torch.from_numpy(y_train)

    outputs=model(inputs)
    loss=criterion(outputs,targets)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch+1) %5==0:
        print("Epoch [{}/{}], LossL:{:.4f}".format(epoch+1,num_epochs,loss.item()))
#计算出训练之后的期望值/预测值，并与实际值进行画图比较
predicted=model(torch.from_numpy(x_train)).detach().numpy()
plt.plot(x_train,y_train,'ro',label='Original Data')
plt.plot(x_train,predicted,label='Fitted Line')
plt.legend()
plt.show()
#保存模型相关数据
torch.save(model.state_dict(),'model.ckpt')

运行结果除了显示一堆文字信息，还有如下图所示的：

 

好了，本文就絮叨到这里了。

 推荐一个公众号：健哥聊量化，会持续推出股票相关基础知识，以及python实现的一些基本的分析代码。欢迎大家关注，二维码如下：