pytorch实现一个简单的回归模型

本文是一个使用pytorch实现简单回归模型的案例，供pytorch初学者参考。

代码整体可以划分为一下几个大的模块：

获取数据、搭建模型、loss、optimizer、训练和测试、保存模型。

导入模块

import torch

获取数据

import numpy as np

import re #正则化模块
ff = open("housing.data").readlines() #逐行读取存储在一个文件内的数据
for item in ff:
    #逐条处理数据,自己根据实际数据保存的特点写
    pass

data = np.array(data).astype(np.float)#转换数据类型

#提取出标签和数据
X = data[:,-1]
Y = data[:,0:-1]

#划分训练集和测试集
X_train = X[0:496,...]
Y_train = Y[0:496,...]
X_test = X[496:,...]
Y_test = Y[496:,...]

搭建网络

class Net(torch.nn.module):
    def __init__(self,n_feature,n_output):#n_feature,n_output分别是特征的数量和输出的数量
        super(Net,self).__init__()
        self.predict = torch.nn.Linear(n_feature,n_output)
        

    def forward(self,x):
        out = self.predict(x)
        return out

#初始化网络
net = Net(13,1)

loss

loss_func = torch.nn.MSELoss()

optimizer

optimizer = torch.optim.SGD(net.parameters(),lr=0.0001)#传入参数和学习率

训练和测试

for i in range(1000):
    x_data = torch.tensor(X_train,dtype=torch.float32)
    y_data = torch.tensor(Y_train,dtype=torch.float32)
    pred = net.forward(x_data)
    #pred和y_data维度不一样，需要处理一下。
    pred = torch.squeeze(pred)
    loss = loss_func(pred, y_data)
    #将梯度置零
    optimizer.zero_grad()
    #反向传播
    loss.backward()
    #更新参数
    optimizer.step()
    #打印loss
    print("ite:{},loss:{}".format(i,loss))
    
    #测试
    x_data = torch.tensor(X_test,dtype=torch.float32)
    y_data = torch.tensor(Y_test,dtype=torch.float32)
    pred = net.forward(x_data)
    loss_test = loss_func(pred, y_data)
    print("ite:{},loss_test:{}".format(i,loss_test))

保存模型

#把模型整体保存,读取时先实例化网络，然后用torch.load()
torch.save(net,"model/model.pkl")
#只保存参数,读取时需要用torch,load_state_dict()
torch.save(net.state_dict()，"model/params.pkl")

这只是一个案例，供学习参考，需要根据自己的数据特点写上相应的数据模块才能运行