Pytorch中（GPU-CPU）实现线性回归

 

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

# 用户选择设备
device_choice = input("请输入您希望使用的设备 (gpu 或 cpu)：").strip().lower()
device = torch.device("cuda" if device_choice == "gpu" and torch.cuda.is_available() else "cpu")
print(f"选择的设备为：{device}")

# 数据生成
np.random.seed(0)
x = np.random.rand(10000, 1)
y = 2 * x + 1 + 0.1 * np.random.randn(10000, 1)
x_train = torch.tensor(x, dtype=torch.float32).to(device)
y_train = torch.tensor(y, dtype=torch.float32).to(device)

# 模型定义
class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(1, 1)  # 一个输入和一个输出

    def forward(self, x):
        return self.linear(x)

model = LinearRegressionModel().to(device)

# 损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 训练过程
losses = []
num_epochs = 200
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    
    # 前向传播
    outputs = model(x_train)
    loss = criterion(outputs, y_train)
    
    # 反向传播和优化
    loss.backward()
    optimizer.step()
    
    losses.append(loss.item())
    
    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 模型参数
w, b = model.linear.weight.item(), model.linear.bias.item()
print(f"模型参数 —— 权重: {w:.4f}, 偏置: {b:.4f}")

# 可视化损失变化
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss Trend')

# 绘制模型预测和实际数据对比
x_test = torch.tensor(np.linspace(0, 1, 100).reshape(-1, 1), dtype=torch.float32).to(device)
with torch.no_grad():
    y_pred = model(x_test).cpu().numpy()

plt.subplot(1, 2, 2)
plt.scatter(x, y, alpha=0.3, label='Actual data')
plt.plot(np.linspace(0, 1, 100), y_pred, color='red', label='Predicted line')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Model Prediction vs Actual Data')
plt.legend()

plt.tight_layout()
plt.show()