以下是完整的操作流程：用 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