在pytorch编写一个神经网络模型,x在0-2范围内拟合指数函数e^x,并更新生成真实曲线和预测曲线之间的动态图
时间: 2024-04-10 22:27:04 浏览: 126
pytorch-神经网络拟合曲线实例
可以使用PyTorch编写一个简单的神经网络模型来拟合指数函数e^x,并使用matplotlib库生成真实曲线和预测曲线之间的动态图。下面是实现的代码示例:
```python
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
# 定义神经网络模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(1, 10)
self.fc2 = nn.Linear(10, 1)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 定义指数函数e^x
def true_function(x):
return np.exp(x)
# 生成训练数据
x_train = np.linspace(0, 2, 100).reshape(-1, 1)
y_train = true_function(x_train)
# 转换为Tensor
x_train = torch.from_numpy(x_train).float()
y_train = torch.from_numpy(y_train).float()
# 初始化神经网络模型和优化器
model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()
# 训练模型
epochs = 1000
losses = []
for epoch in range(epochs):
optimizer.zero_grad()
y_pred = model(x_train)
loss = criterion(y_pred, y_train)
loss.backward()
optimizer.step()
losses.append(loss.item())
# 动态图更新
plt.cla()
plt.scatter(x_train.numpy(), y_train.numpy(), color='b', label='True')
plt.plot(x_train.numpy(), y_pred.detach().numpy(), color='r', label='Predicted')
plt.title('Fitting Exponential Function')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.pause(0.001)
# 绘制损失函数曲线
plt.figure()
plt.plot(range(epochs), losses)
plt.title('Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.show()
```
在上述代码中,我们首先定义了一个简单的神经网络模型`Net`,包含两个全连接层。然后,我们使用指数函数`true_function`生成训练数据,并将其转换为PyTorch的Tensor格式。
接下来,我们初始化模型、优化器和损失函数,并进行模型训练。每个epoch中,我们计算模型的预测值`y_pred`并计算损失函数`loss`,然后通过反向传播更新模型参数。同时,我们使用matplotlib库生成动态图,实时显示真实曲线和预测曲线之间的拟合情况。
最后,我们绘制损失函数随训练epoch的变化曲线,以便观察模型训练过程中的收敛情况。
请注意,由于这是一个简单的示例,可能无法完美拟合指数函数。您可以根据需要调整神经网络模型的结构和超参数以获得更好的拟合效果。
阅读全文