java调用PyTorch 训练模型实现神经网络全流程

以下是完整的操作流程：用 PyTorch 训练模型 → 导出为 ONNX 格式 → 用 Java 加载并推理，兼顾开发效率（PyTorch 快速训练）和生产部署（Java 稳定运行）。

一、PyTorch 训练模型并导出为 ONNX

1. 安装依赖

bash

pip install torch onnx  # PyTorch 和 ONNX 库

2. 训练一个简单模型（以线性回归为例）

python

运行

import torch
import torch.nn as nn
import torch.optim as optim

# 1. 定义模型（线性回归：y = 2x + 3 + 噪声）
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(in_features=1, out_features=1)  # 输入1维，输出1维

    def forward(self, x):
        return self.linear(x)

# 2. 生成训练数据
x_train = torch.tensor([[1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
y_train = torch.tensor([[5.1], [7.2], [8.9], [11.3]], dtype=torch.float32)  # 近似 y=2x+3

# 3. 训练模型
model = LinearModel()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

for epoch in range(1000):
    optimizer.zero_grad()
    y_pred = model(x_train)
    l