用pytorch写一个回归
时间: 2024-05-01 15:21:25 浏览: 77
好的,下面是一个简单的回归模型的示例代码:
```python
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
# 定义模型
class RegressionModel(nn.Module):
def __init__(self):
super(RegressionModel, self).__init__()
self.linear = nn.Linear(1, 1) # 输入为1,输出为1的线性层
def forward(self, x):
y_pred = self.linear(x) # 前向传播
return y_pred
# 准备数据
x_train = np.array([[1.0], [2.0], [3.0], [4.0], [5.0]])
y_train = np.array([[2.0], [4.0], [6.0], [8.0], [10.0]])
x_train = torch.from_numpy(x_train).float()
y_train = torch.from_numpy(y_train).float()
# 初始化模型和优化器
model = RegressionModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 训练模型
num_epochs = 1000
for epoch in range(num_epochs):
# 前向传播
y_pred = model(x_train)
# 计算损失
loss = nn.MSELoss()(y_pred, y_train)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (epoch+1) % 50 == 0:
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
# 测试模型
x_test = torch.tensor([[6.0]])
y_test = model(x_test)
print('The prediction of x=6 is {:.2f}'.format(y_test.item()))
```
这个代码中,我们首先定义了一个简单的线性回归模型,然后准备了一组训练数据和一个优化器。在训练过程中,我们通过前向传播计算模型的预测值,然后计算预测值和真实值之间的均方误差损失,最后通过反向传播和优化器更新模型参数。最后我们用测试数据来测试模型的表现。
阅读全文