一、代码中的数据集可以点击以下链接进行下载

百度网盘提取码：lala

二、代码运行环境

Pytorch-gpu==1.7.1
Python==3.7

三、数据集处理的代码如下所示

import pandas as pd
import torch


# 数据的读取
def make_data():
    data = pd.read_csv(r'dataset\dataset.csv')
    X = data.Education.values.reshape(-1, 1)
    Y = data.Income.values.reshape(-1, 1)
    X = torch.from_numpy(X).type(torch.FloatTensor)
    Y = torch.from_numpy(Y).type(torch.FloatTensor)
    return X, Y


if __name__ == '__main__':
    x, y = make_data()
    print(x.shape)
    print(x.type())

四、模型的构建代码如下所示

from torch import nn


# 模型的构建
class EIModel(nn.Module):
    def __init__(self):
        super(EIModel, self).__init__()
        self.linear = nn.Linear(in_features=1, out_features=1)

    def forward(self, inputs):
        out = self.linear(inputs)
        return out


if __name__ == '__main__':
    model = EIModel()
    print(model)

五、模型的训练代码如下所示

import torch
from data_loader import make_data
from model import EIModel
from torch import nn
from torch import optim
import tqdm

# 数据的读取
X, Y = make_data()

# 模型的构建
model = EIModel()

# 相关的配置
# 损失函数
loss_fn = nn.MSELoss()
# 优化器
opt = optim.SGD(model.parameters(), lr=0.0001)

# 模型的训练
tqdm_range = tqdm.tqdm(range(5000), total=5000)
for epoch in tqdm_range:
    for x, y in zip(X, Y):
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        opt.zero_grad()
        loss.backward()
        opt.step()

# 模型的保存
torch.save(model.state_dict(), r'model_data\model.pth')

六、模型的预测代码如下所示

import torch
from model import EIModel
from data_loader import make_data
import matplotlib.pyplot as plt

# 加载数据
X, Y = make_data()

# 加载模型
model = EIModel()
model_state_dict = torch.load(r'model_data/model.pth')
model.load_state_dict(model_state_dict)

# 模型的验证
# 参数权值的提取
weight, bias = model.linear.weight, model.linear.bias
# 预测效果的展示
plt.scatter(X, Y)
plt.xlabel('Education')
plt.ylabel('Income')
plt.plot(X, model(X).detach().numpy(), c='red')
plt.savefig('result.jpg')
plt.show()