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}")
|
