2025年,许多技术人员仍在疑惑:TensorFlow (TF) 还值得学吗?随着PyTorch在学术研究和新型大模型训练领域的绝对领先,TF似乎正在退居二线。然而,对于专注于AI基础设施和模型部署的资深工程师而言,答案是肯定的——你必须理解TensorFlow的部署生态,尤其是SavedModel格式和TensorFlow Serving (TFS)。
PyTorch解决了“训练什么”的问题,而TensorFlow Serving则解决了“如何可靠且高效地服务数百万用户”的问题。TFS是目前业界最成熟、最稳定、功能最强大的模型服务框架之一,它原生支持SavedModel格式,并提供了批处理、版本控制、A/B测试的API。即使你用PyTorch训练了模型,通过简单的转换,也能利用TFS强大的部署能力。
本文将聚焦核心实操点:如何将一个PyTorch模型通过ONNX过渡到SavedModel,并最终部署到TensorFlow Serving。
步骤一:环境准备和依赖安装
我们需要安装 PyTorch、ONNX 转换工具以及将 ONNX 转换为 SavedModel 的工具。
pip install torch torchvision onnx numpy
pip install onnx-tf # 用于将 ONNX 转换为 TensorFlow SavedModel
步骤二:PyTorch 模型转换为 SavedModel
我们将演示一个简单的图像分类模型的转换流程。
2.1 导出为 ONNX 格式
首先,定义一个简单的PyTorch模型并将其导出为ONNX。
import torch
import torch.nn as nn
import numpy as np
# 1. 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = SimpleModel()
model.eval()
# 2. 准备 dummy 输入
dummy_input = torch.randn(1, 10)
# 3. 导出为 ONNX
onnx_path = "simple_model.onnx"
torch.onnx.export(model,
dummy_input,
onnx_path,
export_params=True,
opset_version=12,
do_constant_folding=True,
input_names=['input'],
output_names=['output'])
print(f"PyTorch model exported to {onnx_path}")
2.2 将 ONNX 转换为 SavedModel
使用 onnx-tf 库将 ONNX 文件转换为 TensorFlow 的 SavedModel 格式。SavedModel是TFS唯一接受的格式,并且必须以版本号(例如1)作为子目录。
from onnx_tf.backend import prepare
import onnx
import os
# 1. 加载 ONNX 模型
onnx_model = onnx.load(onnx_path)
# 2. 转换为 TensorFlow 后端
tf_rep = prepare(onnx_model)
# 3. 定义 SavedModel 的输出路径 (注意:必须包含版本号)
export_dir = "./tf_serving_model/my_pytorch_model/1"
# 4. 导出 SavedModel
os.makedirs(export_dir, exist_ok=True)
tf_rep.export_graph(export_dir)
print(f"SavedModel exported successfully to {export_dir}")
步骤三:使用 TensorFlow Serving 部署模型
SavedModel准备好后,我们使用Docker启动TensorFlow Serving。
# 假设 SavedModel 位于当前目录的 ./tf_serving_model/my_pytorch_model/1
MODEL_BASE_PATH=$(pwd)/tf_serving_model
# 启动 TFS 容器
docker run -t --rm \
-p 8501:8501 \
-v $MODEL_BASE_PATH:/models/ \
tensorflow/serving \
--rest_api_port=8501 \
--model_name=my_pytorch_model \
--model_base_path=/models/my_pytorch_model
# 终端输出将显示服务已启动,监听在 8501 端口。
步骤四:发送推理请求
现在,我们可以通过 RESTful API 向部署在TFS上的“PyTorch模型”发送推理请求。
import requests
import numpy as np
# 1. 准备输入数据 (与 PyTorch 模型的输入维度保持一致: [1, 10])
input_data = np.random.rand(1, 10).astype(np.float32)
# 2. 构建请求 JSON 结构
# 注意 'instances' 字段,TFS 期望 JSON 数组格式
request_data = {
"instances": input_data.tolist()
}
# 3. TFS REST API endpoint
SERVER_URL = 'http://localhost:8501/v1/models/my_pytorch_model:predict'
# 4. 发送 POST 请求
response = requests.post(SERVER_URL, json=request_data)
if response.status_code == 200:
result = response.json()
predictions = np.array(result['predictions'])
print("推理成功!结果形状:", predictions.shape)
print("部分预测结果:", predictions[0][:3])
else:
print(f"请求失败,状态码: {response.status_code}")
print("错误详情:", response.text)
总结:TensorFlow在2025年的价值
对于专注于部署的工程师而言,学习TensorFlow不是学习一个训练框架,而是学习一套工业级的、高可靠性的模型基础设施标准。SavedModel是AI模型在Google生态和许多企业中的“护照”。通过掌握SavedModel的转换和TFS的部署能力,你可以轻松地将任何训练框架(PyTorch、JAX、甚至定制C++模型)的成果,接入到最稳定、性能最优的AI服务管道中。因此,对AI基础设施工程师来说,TensorFlow知识在2025年依然是核心竞争力。
汤不热吧