线性回归模型

线性回归模型

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

# 1.检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# 2.创建模拟数据
np.random.seed(42)
x_train = np.random.rand(100, 1)*10  # 100个样本，1个特征
y_train = 2*x_train + 3 + np.random.randn(100,1) * 2 #y = 2x + 3 + 噪声

# 可视化原始数据
plt.scatter(x_train, y_train, color='blue', label='Original data')
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.show()

# 3.转换为Pytorch张量并移动到GPU
x_train_tensor = torch.from_numpy(x_train).float().to(device)
y_train_tensor = torch.from_numpy(y_train).float().to(device)
# 4.定义线性回归模型
class LinearRegressionModel(nn.Module):
    def __init__(self, input_size, output_size):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_size, output_size)
    def forward(self,x):
        return self.linear(x)
# 5.初始化模型并移动到GPU
input_dim = 1
output_dim = 1

model = LinearRegressionModel(input_dim, output_dim).to(device)

# 6.定义损失函数和优化器
criterion = nn.MSELoss()
learning_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# 7.训练模型
num_epochs = 1000
loss_history = []

for epoch in range(num_epochs):
    #前向传播
    outputs = model(x_train_tensor)
    loss = criterion(outputs, y_train_tensor)

    #反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    #记录损失
    loss_history.append(loss.item())

    #每100个epoch打印一次损失
    if(epoch+1)%100 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 8.可视化损失曲线
plt.plot(loss_history, label='loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss')
plt.show()

# 9. 测试模型
model.eval() # 设置模型为评估模式
with torch.no_grad(): # 禁用梯度计算
    # 生成测试数据
    x_test = np.linspace(0,10,100).reshape(-1,1)
    x_test_tensor = torch.from_numpy(x_test).float().to(device)
    y_pred_tensor = model(x_test_tensor)

    #移回CPU用于绘图
    y_pred = y_pred_tensor.cpu().numpy()

    # 可视化预测结果
    plt.scatter(x_train,y_train, color='blue', label='Original data')
    plt.plot(x_test, y_pred, color='red', label='Fitted line')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.legend()
    plt.show()
    
# 10. 打印学习到的参数
    print("学习到的参数：")
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(f"{name}: {param.data}")

代码说明

1. GPU设置：自动检测并使用GPU（如果可用）
2. 数据生成：创建y=2x+3+噪声的模拟数据
3. 模型定义：简单的单层线性网络
4. 训练过程：
   • 前向传播计算预测值
   • 计算损失（预测值与真实值的差距）
   • 反向传播计算梯度
   • 优化器更新权重
5. 可视化：
   • 原始数据散点图
   • 训练损失曲线
   • 最终拟合直线
6. 参数查看：打印学习到的权重和偏置

如何使用

1. 确保已安装PyTorch（GPU版本）
2. 直接运行代码即可
3. 如果没有GPU，代码会自动使用CPU运行

预期输出

• 开始时Loss较大，随着训练逐渐减小
• 最终拟合的直线应接近y=2x+3
• 打印的参数应接近：weight≈2，bias≈3