如何将PyTorch用作通用函数优化器？求极简无神经网络示例

以下是完全不依赖神经网络层、调度器的极简可运行示例，直接用PyTorch优化器在GPU上拟合自定义参数化函数：

import torch

# 1. 生成模拟数据（自动适配GPU/CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
X = torch.linspace(-5, 5, 1000, device=device)
true_a, true_b, true_c = 2.3, -1.8, 4.2
# 带噪声的真实标签，模拟实际数据
Y = true_a * X**2 + true_b * X + true_c + torch.randn_like(X) * 0.5

# 2. 定义可训练参数（开启梯度追踪，移到对应设备）
a = torch.tensor(1.0, requires_grad=True, device=device)
b = torch.tensor(0.0, requires_grad=True, device=device)
c = torch.tensor(0.0, requires_grad=True, device=device)

# 3. 自定义参数化模型函数
def model(x):
    return a * x**2 + b * x + c

# 4. 损失函数与基础优化器
loss_fn = torch.nn.functional.mse_loss
optimizer = torch.optim.Adam([a, b, c], lr=0.01)

# 5. 训练循环
epochs = 1000
for epoch in range(epochs):
    # 前向传播：计算预测值与损失
    Y_pred = model(X)
    loss = loss_fn(Y_pred, Y)
    
    # 反向传播与参数更新
    optimizer.zero_grad()  # 清空之前的梯度
    loss.backward()        # 计算梯度
    optimizer.step()       # 更新参数
    
    # 每100轮打印训练进度
    if (epoch + 1) % 100 == 0:
        print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}")

# 输出结果对比
print(f"\n优化完成：")
print(f"真实参数 a={true_a:.2f}, b={true_b:.2f}, c={true_c:.2f}")
print(f"拟合参数 a={a.item():.2f}, b={b.item():.2f}, c={c.item():.2f}")

核心要点 #

• GPU适配：所有张量（数据、参数）通过device自动切换GPU/CPU，运算全程在对应设备上执行
• 无神经层依赖：直接用原生张量运算定义参数化模型，完全不需要nn.Linear等神经网络组件
• 极简优化流程：仅使用基础Adam优化器，未引入学习率调度器，核心逻辑和传统numpy/scipy优化一致，只是用PyTorch的自动微分替代手动求导

内容的提问来源于stack exchange，提问作者Domenico