下面是一个使用PyTorch实现全连接网络来拟合简单二次函数
y
=
x
2
y = x^2
y=x2 的示例。我们将创建一个简单的神经网络,定义损失函数和优化器,并进行训练。
下面是完整的代码示例:
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
# 生成数据
x = torch.linspace(-10, 10, 100).unsqueeze(1)
y = x**2
# 定义全连接网络
class SimpleNN(nn.Module):
def __init__(self):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(1, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 1)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
# 实例化模型、损失函数和优化器
model = SimpleNN()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
# 训练模型
num_epochs = 1000
for epoch in range(num_epochs):
model.train()
optimizer.zero_grad()
outputs = model(x)
loss = criterion(outputs, y)
loss.backward()
optimizer.step()
if (epoch+1) % 100 == 0:
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
# 预测
model.eval()
predicted = model(x).detach()
# 可视化结果
plt.plot(x.numpy(), y.numpy(), 'ro', label='Original data')
plt.plot(x.numpy(), predicted.numpy(), 'b-', label='Fitted line')
plt.legend()
plt.show()
这段代码的具体步骤如下:
- 生成数据:创建输入数据 ( x ) 和对应的标签 ( y )。
- 定义网络结构:创建一个简单的全连接神经网络,包括三层线性层。
- 实例化模型、损失函数和优化器:使用均方误差损失函数和Adam优化器。
- 训练模型:在1000个epoch上训练模型,并在每100个epoch打印一次损失。
- 预测和可视化:使用训练好的模型进行预测,并将原始数据和拟合结果进行可视化。
运行这段代码后,你将看到一个图形,其中红点表示原始的二次函数数据,蓝线表示神经网络拟合的结果。