Pytorch实现线性回归并预测结果

news2024/12/24 2:35:06

只有干货, 没有一丝的感情

单输入特征

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

# 设定随机数种子 运行随机函数生成的随机数不会变化
torch.manual_seed(1)

# 学习率
lr = 0.01

# 迭代次数
epochs = 100

# 创建训练集
train_list = [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]

# 创建测试集
test_list = [[11], [12], [13], [14], [15]]

# 将数组转化为tensor
x = torch.tensor(train_list, dtype=torch.float32)
y = torch.tensor(train_list, dtype=torch.float32)
test_x = torch.tensor(test_list, dtype=torch.float32)

# 打印训练集输出
print('训练集输出:', y)

# 定义神经网络各层结构和激活函数
model = nn.Sequential(
    nn.Linear(1, 1),

    # 线性修正单元函数
    nn.ReLU()
)

# 损失函数
cost = nn.MSELoss()

# 迭代优化算法
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 进行训练
for epoch in range(epochs):
    # 前向传播
    outputs = model(x)
    loss = cost(outputs, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch + 1, epochs, loss.item()))

# 预测
test_y = model(test_x)
print('测试集输出:', test_y)

# 画图
# 打开一个窗口
plt.figure()

# 绘制训练集图线
plt.plot(x, y)

# 绘制测试集图线
# tips: 预测结果带梯度,不能直接转为numpy数组,需要忽略梯度才行
plt.plot(test_x, test_y.detach().numpy())

# 把绘制好的图形表示出来
plt.show()

单输入线性回归

多输入特征
以房价为demo

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

# 设定随机数种子 运行随机函数生成的随机数不会变化
torch.manual_seed(1)

# 学习率
lr = 0.000003

# 迭代次数
epochs = 100

# 创建训练集
# 不同列分别代表含义 房屋面积 客厅+卧室数量 客厅面积
train_input = [[50, 2, 10], [60, 2, 15], [69, 3, 12], [75, 3, 12], [80, 3, 18],
               [89, 4, 15], [97, 4, 20], [109, 4, 27], [129, 5, 20], [150, 5, 25]]

# 输出 售价(万元)
train_output = [[100], [115], [140], [149], [165],
                [195], [215], [235], [260], [310]]
# 创建测试集
test_list = [[33, 2, 7], [55, 2, 12], [67, 2, 15], [88, 3, 20], [99, 3, 20]]

# 将数组转化为tensor
x = torch.tensor(train_input, dtype=torch.float32)
y = torch.tensor(train_output, dtype=torch.float32)
test_x = torch.tensor(test_list, dtype=torch.float32)

# 打印训练集输出
print('训练集输出:', y)

# 定义神经网络各层结构和激活函数
model = nn.Sequential(
    nn.Linear(x.size()[1], 1)
)

# 损失函数
cost = nn.MSELoss()

# 迭代优化算法
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 进行训练
for epoch in range(epochs):
    # 前向传播
    outputs = model(x)
    loss = cost(outputs, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch + 1, epochs, loss.item()))

# 预测
test_y = model(test_x)
print('测试集输出:', test_y)

# 画图
# 打开一个窗口
plt.figure()

# 绘制训练集图线
plt.plot(x, y, label='feature')

# 绘制测试集图线
# tips: 预测结果带梯度,不能直接转为numpy数组,需要忽略梯度才行
plt.plot(test_x, test_y.detach().numpy(), label='predict')
plt.legend()

# 把绘制好的图形表示出来
plt.show()