在 Hugging Face 中使用 TF Serving 部署 TensorFlow 视觉模型

为 TF Serving 保存模型

要使用 TF Serving 部署模型，模型必须采用 SavedModel 格式。Hugging Face Transformers 库中的 TensorFlow 模型提供了 save_pretrained() 方法，可将权重序列化为 h5 和 SavedModel 两种格式。

对于 Vision Transformer（ViT）模型，流程包括通过 TFViTForImageClassification.from_pretrained() 加载模型，然后调用 save_pretrained(saved_model=True)。默认情况下，这会创建 TF Serving 所需的带版本号的目录结构（例如 {model_dir}/saved_model/{version}）。

实现模型手术以进行前后处理

标准的机器学习模型需要特定的前处理和后处理。为降低训练‑服务偏差并减轻开发者的认知负担，这些操作可以通过“模型手术”直接嵌入模型的计算图中。

前处理流水线

对于 ViT 模型，关键的前处理步骤包括：

  • Scaling：将图像像素值转换到 [0, 1] 范围。
  • Normalization：使用模型特定的均值和标准差，将像素值缩放到 [-1, 1] 范围。
  • Resizing：将图像调整为 224×224 的空间分辨率。
  • Transposition：将通道维度移动到前面，以匹配 Hugging Face 模型使用的 channel‑first 格式。

为优化请求负载并防止体积膨胀，指南建议接受 base64 编码的字符串作为输入，然后在计算图中使用 tf.io.decode_base64 和 tf.io.decode_jpeg 进行解码和处理。

后处理与导出

后处理将原始模型 logits 转换为人类可读的标签。通过从模型的 call() 方法派生具体函数，开发者可以将模型包装在 serving_fn 中，该函数：

  1. 执行前处理流水线。
  2. 进行模型推理。
  3. 对 logits 应用 softmax 函数，以得到置信分数。
  4. 使用模型的 id2label 配置，将得到的索引映射为字符串标签。

随后使用 tf.saved_model.save() 将该包装函数导出为 serving_default 签名，将模型的输入要求从 4D 张量改为字符串，并将输出改为包含标签和置信分数的字典。

使用 TensorFlow Serving 部署

模型导出后，可使用 tensorflow_model_server 命令进行部署。关键配置参数包括：

  • rest_api_port：REST 接口的端口（默认通常为 8501）。
  • model_name：调用 API 时使用的标识符。
  • model_base_path：TF Serving 加载最新模型版本的目录。

TF Serving 提供两种主要方式来查询已部署的模型：

REST 接口

REST API 适用于在线预测场景。请求以 JSON 负载形式发送，其中 instances 列表包含 base64 编码的图像。接口遵循以下格式：http://localhost:8501/v1/models/{model_name}:predict。

gRPC 接口

对于低延迟、高可扩展性和分布式系统，gRPC 是首选的部署方式。这需要通过 grpc.insecure_channel 打开通信通道，并使用 PredictionServiceStub 发送 PredictRequest 负载。gRPC 返回包含置信分数和预测标签的结构化输出。

Sources