線性迴歸的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實現的一些基本的分析代碼。歡迎大家關注，二維碼如下：