pytorch线性回归

最新推荐文章于 2024-09-10 23:52:43 发布

小李小李快乐w比

最新推荐文章于 2024-09-10 23:52:43 发布

阅读量107

点赞数

文章标签： pytorch 线性回归 python

本文链接：https://blog.csdn.net/weixin_43480152/article/details/130875667

版权

该代码示例展示了如何在PyTorch中构建一个简单的线性回归模型，使用均方误差（MSELoss）作为损失函数和随机梯度下降（SGD）优化器进行训练。数据集是手动定义的一组坐标点，模型经过训练后能拟合这些点并生成一条最佳拟合直线。

摘要由CSDN通过智能技术生成

线性回归

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

# 超参数 输入输出的维度、训练轮数和学习率
input_size = 1
output_size = 1
num_epochs = 100
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)

# 定义线性回归模型 类似于 y=w*x+b
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):
    # 将输入数据转化为pytorch张量
    inputs = torch.from_numpy(x_train)
    targets = torch.from_numpy(y_train)

    # 前向传播，将输入数据传递给模型，计算输出结果outputs和目标标签targets之间的损失值，使用定义的损失函数计算MSE损失
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # 梯度缓存清0，执行反向传播计算，使用优化器更新模型参数。
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

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

# plot the graph
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()

# Save the model checkpoint
torch.save(model.state_dict(), 'model.ckpt')