import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
# 1. 创建线性回归模型类
class LinearRegressionModel(nn.Module):
def __init__(self):
super(LinearRegressionModel, self).__init__()
self.linear = nn.Linear(1, 1) # 1个输入特征,1个输出
def forward(self, x):
return self.linear(x)
# 2. 生成训练数据
area = np.array([1000, 1500, 1800, 2400, 3000], dtype=np.float32).reshape(-1, 1) # 房屋面积(平方英尺)
price = np.array([250000, 300000, 350000, 500000, 600000], dtype=np.float32).reshape(-1, 1) # 房价
# 标准化房屋面积
area = area / 3000 # 假设最大面积为3000平方英尺
# 转换为 PyTorch 张量
x_train = torch.from_numpy(area)
y_train = torch.from_numpy(price)
# 3. 实例化模型、定义损失函数和优化器
model = LinearRegressionModel()
criterion = nn.MSELoss() # 均方误差损失函数
optimizer = optim.SGD(model.parameters(), lr=0.001) # 学习率调低
# 4. 训练模型
epochs = 1000
for epoch in range(epochs):
# 前向传播
outputs = model(x_train)
loss = criterion(outputs, y_train)
# 反向传播
optimizer.zero_grad() # 清零梯度
loss.backward() # 计算梯度
optimizer.step() # 更新权重
# 每100次输出一次损失值
if (epoch + 1) % 100 == 0:
print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')
# 5. 保存训练好的模型
torch.save(model.state_dict(), 'linear_regression_model.pth')
print("模型已保存!")
# 6. 加载模型并进行预测
loaded_model = LinearRegressionModel()
loaded_model.load_state_dict(torch.load('linear_regression_model.pth'))
loaded_model.eval() # 设置为评估模式
# 进行预测
new_area = torch.tensor([[2500 / 3000]], dtype=torch.float32) # 假设新房屋面积为2500平方英尺,标准化处理
predicted_price = loaded_model(new_area)
print(f"Predicted price for area 2500 sq.ft: ${predicted_price.item():,.2f}")
- 创建模型:
LinearRegressionModel
是一个简单的线性回归模型,只有一个线性层 (nn.Linear
)。 - 数据准备:我们生成了一个简单的示例数据集,包含房屋面积和房价数据。数据被转换为 PyTorch 张量格式。
- 模型训练:使用均方误差损失函数 (
MSELoss
) 和随机梯度下降优化器 (SGD
) 来训练模型。模型在1000个迭代中进行训练,并在每100次迭代后输出损失值。 - 保存模型:训练完成后,使用
torch.save
保存模型的参数。 - 加载模型并进行预测:使用
torch.load
加载模型参数,并将模型设置为评估模式 (eval
)。然后,我们通过模型对一个新的房屋面积值进行预测。