汤不热吧在AI基础设施领域，一个常见的问题是：NVIDIA在内部究竟更倾向于使用PyTorch还是TensorFlow？答案是，NVIDIA作为硬件和基础设施提供商，其核心目标是确保所有主流框架都能在其CUDA平台上高效运行。因此，NVIDIA的策略是框架中立，但鼓励使用能最大化硬件性能的工具。
NVIDIA通过其软件栈（如CUDA, cuDNN, Triton Inference Server）同时支持PyTorch和TensorFlow。然而，在模型部署阶段，为了实现统一、高效、可移植的推理，业界和NVIDIA自身都大力推崇使用ONNX (Open Neural Network Exchange) 标准。
本文将聚焦如何利用ONNX Runtime，在NVIDIA GPU上构建一个统一的推理管道，从而彻底解决框架选择带来的部署复杂性。
1. 为什么选择ONNX进行部署？
ONNX是一个开放标准，定义了一套可扩展的计算图格式。它允许模型从一个框架（如PyTorch或TensorFlow）训练完成后，转换成ONNX格式，然后在另一个框架或运行时（如ONNX Runtime, 尤其在GPU上使用onnxruntime-gpu）中执行。
关键优势：
1. 性能最大化： ONNX Runtime支持多种硬件加速器，在NVIDIA GPU上，它可以直接利用TensorRT或CUDA执行内核，通常比原生框架的推理性能更高。
2. 框架解耦： 部署团队无需关心模型是用哪个框架训练的。
3. Triton集成： NVIDIA的Triton Inference Server原生支持ONNX Runtime后端，使其成为高效生产部署的首选。
2. 准备工作：安装必要的库
为了演示，我们需要安装PyTorch、TensorFlow和ONNX相关的库。请确保您的系统已安装CUDA驱动，以便使用onnxruntime-gpu。
pip install torch torchvision tensorflow onnx onnxruntime-gpu
pip install tf2onnx # 用于TensorFlow模型转换
3. 统一ONNX：从PyTorch和TensorFlow导出
我们将展示如何将一个简单的线性/全连接模型导出为ONNX格式。
3.1 导出 PyTorch 模型到 ONNX
PyTorch通过内置的torch.onnx.export函数支持直接导出。
import torch
import torch.nn as nn
# 1. 定义一个简单的PyTorch模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.linear = nn.Linear(10, 5)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# 2. 准备一个Dummy输入（必须）
dummy_input = torch.randn(1, 10)
# 3. 导出模型
onnx_path_pt = "pytorch_model.onnx"
torch.onnx.export(
model,
dummy_input,
onnx_path_pt,
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}}
)
print(f"PyTorch模型已导出到: {onnx_path_pt}")
3.2 导出 TensorFlow 模型到 ONNX
TensorFlow模型转换通常需要依赖tf2onnx库。
import tensorflow as tf
import tf2onnx
# 1. 定义一个简单的TensorFlow模型
model_tf = tf.keras.Sequential([
tf.keras.layers.Dense(5, input_shape=(10,))
])
# 2. 准备签名和Dummy输入
input_signature = [
tf.TensorSpec([None, 10], tf.float32, name='input')
]
# 3. 导出模型
onnx_path_tf = "tensorflow_model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
model_tf,
input_signature=input_signature,
opset=13,
output_path=onnx_path_tf
)
print(f"TensorFlow模型已导出到: {onnx_path_tf}")
4. 使用ONNX Runtime在GPU上统一推理
现在我们有两个ONNX文件，它们都遵循相同的标准。我们可以使用同一套ONNX Runtime代码，并强制它使用GPU加速（通过CUDAProvider），实现高性能的统一推理。
import onnxruntime
import numpy as np
# 准备测试数据
test_input = np.random.rand(1, 10).astype(np.float32)
# 定义要测试的模型列表
onnx_files = ["pytorch_model.onnx", "tensorflow_model.onnx"]
for model_path in onnx_files:
print(f"\n--- 推理模型: {model_path} ---")
# 1. 创建ONNX Runtime Session
# 明确指定providers: 'CUDAExecutionProvider' 优先于 'CPUExecutionProvider'
try:
sess = onnxruntime.InferenceSession(
model_path,
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)
except Exception as e:
print(f"警告: CUDA提供者加载失败 ({e})。尝试回退到CPU。")
sess = onnxruntime.InferenceSession(
model_path,
providers=['CPUExecutionProvider']
)
# 2. 获取输入输出名称
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
# 3. 执行推理
input_data = {input_name: test_input}
result = sess.run([output_name], input_data)
print("推理结果形状:", result[0].shape)
print("推理结果 (前5个元素):")
print(result[0][0, :5])
# 4. 验证执行提供者 (确认是否使用了GPU)
print("当前使用的执行提供者:", sess.get_providers()[0])
总结
NVIDIA的生态系统致力于提供最高效的加速层（CUDA, TensorRT）。无论模型是在PyTorch还是TensorFlow中训练，通过将模型统一到ONNX格式，我们能确保在NVIDIA GPU上运行时能够充分利用这些底层加速技术，实现部署的标准化和性能的最优化。在生产环境中，这正是NVIDIA基础设施所推荐和优化的部署路径。