使用Pytorch用神经网络拟合高斯曲线

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

# 生成高斯曲线数据
mean = 0
std = 1
x_train = np.linspace(-5, 5, 100)
y_train = np.exp(-np.power(x_train - mean, 2) / (2 * np.power(std, 2)))

class GaussianNet(nn.Module):
    def __init__(self):
        super(GaussianNet, 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

model = GaussianNet()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

num_epochs = 1000
for epoch in range(num_epochs):
    inputs = torch.from_numpy(x_train).float().view(-1, 1)
    labels = torch.from_numpy(y_train).float().view(-1, 1)

    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

x_test = np.linspace(-5, 5, 100)
inputs = torch.from_numpy(x_test).float().view(-1, 1)
outputs = model(inputs).detach().numpy()

plt.plot(x_train, y_train, label='Ground Truth')
plt.plot(x_test, outputs, label='Fitted Curve')
plt.legend()
plt.show()