介绍

下面是基于pytorch的神经网络基本框架代码，这里的任务是让神经网络拟合sin函数

import torch
import torch.nn as nn

# ===== 0. 造数据：拟合 y = sin(x) 这个非线性函数 =====
torch.manual_seed(0)
x = torch.linspace(-3.14, 3.14, 200).unsqueeze(1)   # [200, 1]
y = torch.sin(x)                                     # [200, 1]

# 1.定义模型(万能骨架)
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1,32)
        self.fc2 = nn.Linear(32,32)
        self.fc3 = nn.Linear(32,1)
    def forward(self,x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

model = MLP()

# ===== 2. 损失函数 + 优化器 =====
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 3. 训练循环，经典5步走
for epoch in range(2000):
    optimizer.zero_grad()
    pred = model(x)
    loss = criterion(pred, y)
    loss.backward()
    optimizer.step()
    if epoch%100==0:
        print(f'当前第{epoch}轮，loss为{loss.item():.6f}')

# 4. 检验学习结果
test_x = torch.tensor([[0.0], [1.5708], [-1.5708]]) # 0,π/2, -π/2
test_pred = model(test_x)
print("预测 sin(0)   =", round(test_pred[0].item(), 4), " 真实 =", 0.0)
print("预测 sin(π/2) =", round(test_pred[1].item(), 4), " 真实 =", 1.0)
print("预测 sin(-π/2)=", round(test_pred[2].item(), 4), " 真实 =", -1.0)