PyTorch实战：全连接网络拟合简单函数

PyTorch实战：全连接网络拟合简单函数

使用神经网络拟合简单函数是理解神经网络工作原理的一个很好的起点。本文将通过一个具体的实例，展示如何使用PyTorch实现全连接网络来拟合二次函数y = x^2。

代码实现

下面是完整的代码示例：

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()

代码解析

这段代码的具体步骤如下：

1. 生成数据：创建输入数据 ( x ) 和对应的标签 ( y )。
2. 定义网络结构：创建一个简单的全连接神经网络，包括三层线性层。
3. 实例化模型、损失函数和优化器：使用均方误差损失函数和Adam优化器。
4. 训练模型：在1000个epoch上训练模型，并在每100个epoch打印一次损失。
5. 预测和可视化：使用训练好的模型进行预测，并将原始数据和拟合结果进行可视化。

运行这段代码后，你将看到一个图形，其中红点表示原始的二次函数数据，蓝线表示神经网络拟合的结果。