如何快速上手Forge Pump Surrogate：从ONNX到TensorFlow的多运行时推理教程-平芜编程栈

如何快速上手Forge Pump Surrogate：从ONNX到TensorFlow的多运行时推理教程

【免费下载链接】forge-pump-surrogate-multiruntime项目地址: https://ai.gitcode.com/hf_mirrors/sankalpsthakur/forge-pump-surrogate-multiruntime

Forge Pump Surrogate是一个强大的泵代理模型工具，支持ONNX、PyTorch和TensorFlow等多种运行时环境，能帮助开发者轻松实现跨平台的泵系统推理。本教程将带你快速掌握从环境搭建到多运行时推理的完整流程，让你在不同场景下都能高效使用这个工具。

📋 环境准备：轻松搭建开发环境

要开始使用Forge Pump Surrogate，首先需要准备好开发环境。你需要安装Python以及相关的依赖库。项目的依赖信息在requirements.txt文件中，里面列出了所有必要的Python包及其版本。

你可以使用以下命令克隆项目仓库并安装依赖：

git clone https://gitcode.com/hf_mirrors/sankalpsthakur/forge-pump-surrogate-multiruntime cd forge-pump-surrogate-multiruntime pip install -r requirements.txt

安装完成后，你就拥有了使用Forge Pump Surrogate的基本环境。这个环境支持后续的模型训练、转换和推理等所有操作。

🔍 项目结构解析：了解核心组件

Forge Pump Surrogate的项目结构清晰，各个目录和文件都有明确的功能划分。让我们来了解一下主要的组成部分：

edge/：包含边缘设备相关的推理代码，如edge/inference_onnx.py是ONNX格式模型的推理脚本。
onnx/：存放ONNX格式的模型文件onnx/model.onnx。
pytorch/：包含PyTorch相关的模型文件，如model_state.pt是模型的状态文件。
src/：源代码目录，其中src/model.py定义了PumpSurrogate模型，src/build_release.py是构建和导出模型的关键脚本。
tensorflow/：存放TensorFlow格式的模型，如model.keras和model.tflite。

通过这个结构，你可以很容易地找到不同运行时环境下的模型和相关代码，为后续的使用和扩展提供了便利。

🚀 模型导出流程：从PyTorch到多运行时

Forge Pump Surrogate的一大特色是支持多种运行时环境，这得益于其完善的模型导出流程。在src/build_release.py中，定义了将PyTorch模型导出为ONNX、TensorFlow和TFLite格式的完整过程。

PyTorch模型保存

首先，训练好的PyTorch模型会被保存为状态文件和脚本文件：

torch.save(model.state_dict(), model_dir / "pytorch" / "model_state.pt") traced = torch.jit.trace(model, example) traced.save(str(model_dir / "pytorch" / "model.ts"))

导出为ONNX格式

接着，使用PyTorch的ONNX导出功能将模型转换为ONNX格式：

onnx_path = model_dir / "onnx" / "model.onnx" torch.onnx.export( model, example, onnx_path, input_names=["features"], output_names=["outputs"], dynamic_axes={"features": {0: "batch"}, "outputs": {0: "batch"}}, opset_version=18, dynamo=False, )

导出为TensorFlow和TFLite格式

最后，通过自定义的export_tensorflow函数将模型转换为TensorFlow的Keras格式和TFLite格式：

def export_tensorflow(model: PumpSurrogate, model_dir: Path) -> tuple[Path, Path]: # ... 代码省略 ... keras_path = model_dir / "tensorflow" / "model.keras" keras_model.save(keras_path) converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) tflite = converter.convert() tflite_path = model_dir / "tensorflow" / "model.tflite" tflite_path.write_bytes(tflite) return keras_path, tflite_path

这个完整的导出流程确保了模型可以在不同的运行时环境中使用，极大地扩展了模型的应用场景。

💻 多运行时推理教程：轻松实现跨平台部署

Forge Pump Surrogate支持在多种运行时环境下进行推理，下面我们分别介绍在ONNX、PyTorch和TensorFlow环境下的推理方法。

ONNX运行时推理

ONNX格式的模型可以使用ONNX Runtime进行推理。在edge/inference_onnx.py中，定义了使用ONNX模型进行推理的方法：

parser.add_argument("--model", type=Path, default=Path("onnx/model.onnx")) # ... 代码省略 ... ort_session = ort.InferenceSession(str(args.model), providers=["CPUExecutionProvider"]) results = ort_session.run(["outputs"], {"features": input_data.astype(np.float32)})

PyTorch运行时推理

PyTorch模型可以直接加载进行推理：

model = PumpSurrogate(input_mean, input_std, output_mean, output_std) model.load_state_dict(torch.load("pytorch/model_state.pt")) model.eval() with torch.no_grad(): prediction = model(input_tensor)

TensorFlow运行时推理

TensorFlow的Keras模型和TFLite模型也都可以方便地进行推理：

# Keras模型推理 keras_model = tf.keras.models.load_model("tensorflow/model.keras") prediction = keras_model(input_data, training=False).numpy() # TFLite模型推理 interpreter = tf.lite.Interpreter(model_path="tensorflow/model.tflite") interpreter.allocate_tensors() input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() interpreter.set_tensor(input_details[0]['index'], input_data) interpreter.invoke() prediction = interpreter.get_tensor(output_details[0]['index'])

通过这些简单的代码，你可以在不同的运行时环境中轻松实现模型推理，满足各种部署需求。

📊 模型性能评估：确保推理质量

Forge Pump Surrogate还提供了完善的模型性能评估功能。在src/build_release.py中，定义了回归指标计算和跨运行时一致性检查等功能。

回归指标计算

通过regression_metrics函数可以计算模型的MAE、RMSE和R²等指标：

def regression_metrics(actual: np.ndarray, predicted: np.ndarray) -> dict[str, dict[str, float]]: metrics: dict[str, dict[str, float]] = {} for index, name in enumerate(TARGET_NAMES): truth = actual[:, index] pred = predicted[:, index] mae = float(np.mean(np.abs(pred - truth))) rmse = float(np.sqrt(np.mean((pred - truth) ** 2))) denom = float(np.sum((truth - np.mean(truth)) ** 2)) r2 = 1.0 - float(np.sum((truth - pred) ** 2)) / denom if denom else 1.0 metrics[name] = {"mae": mae, "rmse": rmse, "r2": r2} return metrics

跨运行时一致性检查

为了确保不同运行时环境下模型推理结果的一致性，项目中还进行了严格的一致性检查：

onnx_delta = float(np.max(np.abs(torch_pred[:32] - onnx_pred[:32]))) tensorflow_delta = float(np.max(np.abs(torch_pred[:32] - keras_pred[:32]))) litert_delta = float(np.max(np.abs(torch_pred[:32] - tflite_pred))) conformance_passed = max(onnx_delta, tensorflow_delta, litert_delta) <= 1e-3

这些评估功能确保了你使用的模型具有良好的性能和跨平台一致性，可以放心地在各种场景中应用。