在 Hugging Face 中使用 TF Serving 部署 TensorFlow 视觉模型
在过去几个月里，Hugging Face 团队和外部贡献者在 Transformers 中为 TensorFlow 添加了各种视觉模型。这个列表正在全面增长，并且已经包含了最先进的预训练模型，如 Vision Transformer、Masked Autoencoders、RegNet、ConvNeXt 等等！
在部署 TensorFlow 模型时，您有多种选择。根据您的用例，您可能希望将模型公开为端点或将其打包到应用程序本身。TensorFlow 提供了适用于这些不同场景的工具。
在这篇文章中，您将看到如何使用 TensorFlow Serving (TF Serving) 在本地部署 Vision Transformer (ViT) 模型（用于图像分类）。这将允许开发人员将模型公开为 REST 或 gRPC 端点。此外，TF Serving 原生支持许多部署特定功能，如模型预热、服务器端批处理等。
要获取本文中显示的所有完整工作代码，请参阅开头的 Colab Notebook。
保存模型
所有 🤗 Transformers 中的 TensorFlow 模型都有一个名为 save_pretrained() 的方法。通过它，您可以将模型权重序列化为 h5 格式以及独立的 SavedModel 格式。TF Serving 需要模型以 SavedModel 格式存在。因此，我们首先加载一个 Vision Transformer 模型并保存它。
from transformers import TFViTForImageClassification
temp_model_dir = "vit"
ckpt = "google/vit-base-patch16-224"
model = TFViTForImageClassification.from_pretrained(ckpt)
model.save_pretrained(temp_model_dir, saved_model=True)
默认情况下，save_pretrained() 将首先在我们提供的路径中创建一个版本目录。因此，最终路径变为：{temp_model_dir}/saved_model/{version}。
我们可以像这样检查 SavedModel 的服务签名
saved_model_cli show --dir {temp_model_dir}/saved_model/1 --tag_set serve --signature_def serving_default
这应该输出
The given SavedModel SignatureDef contains the following input(s):
inputs['pixel_values'] tensor_info:
dtype: DT_FLOAT
shape: (-1, -1, -1, -1)
name: serving_default_pixel_values:0
The given SavedModel SignatureDef contains the following output(s):
outputs['logits'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1000)
name: StatefulPartitionedCall:0
Method name is: tensorflow/serving/predict
可以看出，模型接受单个 4 维输入（即 pixel_values），其轴为：(batch_size, num_channels, height, width)。对于此模型，可接受的高度和宽度设置为 224，通道数为 3。您可以通过检查模型的 config 参数 (model.config) 来验证这一点。模型生成一个 1000 维的 logits 向量。
模型改造
通常，每个机器学习模型都有特定的预处理和后处理步骤。ViT 模型也不例外。主要的预处理步骤包括：
将图像像素值缩放到 [0, 1] 范围。
将缩放后的像素值归一化到 [-1, 1]。
调整图像大小，使其空间分辨率为 (224, 224)。
您可以通过检查与模型关联的图像处理器来确认这些内容
from transformers import AutoImageProcessor
processor = AutoImageProcessor.from_pretrained(ckpt)
print(processor)
这应该打印
ViTImageProcessor {
"do_normalize": true,
"do_resize": true,
"image_mean": [
0.5,
0.5,
0.5
],
"image_std": [
0.5,
0.5,
0.5
],
"resample": 2,
"size": 224
}
由于这是一个在 ImageNet-1k 数据集上预训练的图像分类模型，因此模型的输出需要作为后处理步骤映射到 ImageNet-1k 类。
为了减少开发人员的认知负荷和训练-服务偏差，通常最好发布一个内置了大部分预处理和后处理步骤的模型。因此，您应该将模型序列化为 SavedModel，以便将上述处理操作嵌入到其计算图中。
预处理
对于预处理，图像归一化是最重要的组成部分之一
def normalize_img(
img, mean=processor.image_mean, std=processor.image_std
):
# Scale to the value range of [0, 1] first and then normalize.
img = img / 255
mean = tf.constant(mean)
std = tf.constant(std)
return (img - mean) / std
您还需要调整图像大小并进行转置，以便它具有领先的通道维度，因为这遵循了 🤗 Transformers 的标准格式。以下代码片段显示了所有预处理步骤
CONCRETE_INPUT = "pixel_values" # Which is what we investigated via the SavedModel CLI.
SIZE = processor.size["height"]
def normalize_img(
img, mean=processor.image_mean, std=processor.image_std
):
# Scale to the value range of [0, 1] first and then normalize.
img = img / 255
mean = tf.constant(mean)
std = tf.constant(std)
return (img - mean) / std
def preprocess(string_input):
decoded_input = tf.io.decode_base64(string_input)
decoded = tf.io.decode_jpeg(decoded_input, channels=3)
resized = tf.image.resize(decoded, size=(SIZE, SIZE))
normalized = normalize_img(resized)
normalized = tf.transpose(
normalized, (2, 0, 1)
) # Since HF models are channel-first.
return normalized
@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def preprocess_fn(string_input):
decoded_images = tf.map_fn(
preprocess, string_input, dtype=tf.float32, back_prop=False
)
return {CONCRETE_INPUT: decoded_images}
关于使模型接受字符串输入的说明:
通过 REST 或 gRPC 请求处理图像时，请求有效负载的大小可能会根据传递图像的分辨率而轻易增加。这就是为什么可靠地压缩它们然后准备请求有效负载是一个很好的做法。
后处理和模型导出
您现在已经掌握了可以注入到模型现有计算图中的预处理操作。在本节中，您还将后处理操作注入到图中并导出模型！
def model_exporter(model: tf.keras.Model):
m_call = tf.function(model.call).get_concrete_function(
tf.TensorSpec(
shape=[None, 3, SIZE, SIZE], dtype=tf.float32, name=CONCRETE_INPUT
)
)
@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def serving_fn(string_input):
labels = tf.constant(list(model.config.id2label.values()), dtype=tf.string)
images = preprocess_fn(string_input)
predictions = m_call(**images)
indices = tf.argmax(predictions.logits, axis=1)
pred_source = tf.gather(params=labels, indices=indices)
probs = tf.nn.softmax(predictions.logits, axis=1)
pred_confidence = tf.reduce_max(probs, axis=1)
return {"label": pred_source, "confidence": pred_confidence}
return serving_fn
您可以首先从模型的前向传递方法 (call()) 中导出具体函数，以便模型能够很好地编译成图。之后，您可以按顺序执行以下步骤
通过预处理操作传递输入。
通过导出的具体函数传递预处理输入。
对输出进行后处理，并以格式良好的字典形式返回它们。
现在是导出模型的时候了！
MODEL_DIR = tempfile.gettempdir()
VERSION = 1
tf.saved_model.save(
model,
os.path.join(MODEL_DIR, str(VERSION)),
signatures={"serving_default": model_exporter(model)},
)
os.environ["MODEL_DIR"] = MODEL_DIR
导出后，我们再次检查模型签名
saved_model_cli show --dir {MODEL_DIR}/1 --tag_set serve --signature_def serving_default
The given SavedModel SignatureDef contains the following input(s):
inputs['string_input'] tensor_info:
dtype: DT_STRING
shape: (-1)
name: serving_default_string_input:0
The given SavedModel SignatureDef contains the following output(s):
outputs['confidence'] tensor_info:
dtype: DT_FLOAT
shape: (-1)
name: StatefulPartitionedCall:0
outputs['label'] tensor_info:
dtype: DT_STRING
shape: (-1)
name: StatefulPartitionedCall:1
Method name is: tensorflow/serving/predict
您会注意到模型的签名现在已经改变了。具体来说，输入类型现在是字符串，模型返回两项内容：置信度分数和字符串标签。
如果您已经安装了 TF Serving（Colab Notebook 中有介绍），那么您现在就可以部署此模型了！
使用 TensorFlow Serving 部署
只需一个命令即可完成此操作
nohup tensorflow_model_server \
--rest_api_port=8501 \
--model_name=vit \
--model_base_path=$MODEL_DIR >server.log 2>&1
从上述命令中，重要的参数是
rest_api_port表示 TF Serving 将用于部署模型 REST 端点的端口号。默认情况下，TF Serving 使用 8500 端口作为 gRPC 端点。model_name指定将用于调用 API 的模型名称（可以是任何名称）。model_base_path表示 TF Serving 将用于加载模型最新版本的模型基础路径。
（支持参数的完整列表在此处。）
瞧！几分钟之内，您就可以部署一个拥有 REST 和 gRPC 两个端点的模型了。
查询 REST 端点
回想一下，您导出的模型接受用 base64 格式编码的字符串输入。因此，要构建请求有效负载，您可以这样做
# Get image of a cute cat.
image_path = tf.keras.utils.get_file(
"image.jpg", "http://images.cocodataset.org/val2017/000000039769.jpg"
)
# Read the image from disk as raw bytes and then encode it.
bytes_inputs = tf.io.read_file(image_path)
b64str = base64.urlsafe_b64encode(bytes_inputs.numpy()).decode("utf-8")
# Create the request payload.
data = json.dumps({"signature_name": "serving_default", "instances": [b64str]})
TF Serving 的 REST 端点请求负载格式规范可在此处获取。在 instances 中，您可以传递多个编码图像。此类端点旨在用于在线预测场景。对于具有多个数据点的输入，您可能需要启用批处理以获得性能优化优势。
现在您可以调用 API
headers = {"content-type": "application/json"}
json_response = requests.post(
"https://:8501/v1/models/vit:predict", data=data, headers=headers
)
print(json.loads(json_response.text))
# {'predictions': [{'label': 'Egyptian cat', 'confidence': 0.896659195}]}
REST API 为 - https://:8501/v1/models/vit:predict，遵循此处的规范。默认情况下，它总是获取模型的最新版本。但如果您想要特定版本，可以这样做：https://:8501/v1/models/vit/versions/1:predict。
查询 gRPC 端点
虽然 REST 在 API 领域非常流行，但许多应用程序通常受益于 gRPC。这篇文章很好地比较了两种部署方式。gRPC 通常更适用于低延迟、高度可伸缩的分布式系统。
有几个步骤。首先，您需要打开一个通信通道
import grpc
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
channel = grpc.insecure_channel("localhost:8500")
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
然后，创建请求有效负载
request = predict_pb2.PredictRequest()
request.model_spec.name = "vit"
request.model_spec.signature_name = "serving_default"
request.inputs[serving_input].CopyFrom(tf.make_tensor_proto([b64str]))
您可以像这样通过编程方式确定 serving_input 键
loaded = tf.saved_model.load(f"{MODEL_DIR}/{VERSION}")
serving_input = list(
loaded.signatures["serving_default"].structured_input_signature[1].keys()
)[0]
print("Serving function input:", serving_input)
# Serving function input: string_input
现在，您可以获得一些预测结果
grpc_predictions = stub.Predict(request, 10.0) # 10 secs timeout
print(grpc_predictions)
outputs {
key: "confidence"
value {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 1
}
}
float_val: 0.8966591954231262
}
}
outputs {
key: "label"
value {
dtype: DT_STRING
tensor_shape {
dim {
size: 1
}
}
string_val: "Egyptian cat"
}
}
model_spec {
name: "resnet"
version {
value: 1
}
signature_name: "serving_default"
}
您还可以从上述结果中获取感兴趣的键值，如下所示
grpc_predictions.outputs["label"].string_val, grpc_predictions.outputs[
"confidence"
].float_val
# ([b'Egyptian cat'], [0.8966591954231262])
总结
在这篇文章中，我们学习了如何使用 TF Serving 部署 Transformers 中的 TensorFlow 视觉模型。虽然本地部署对于周末项目很棒，但我们希望能够扩展这些部署以服务许多用户。在接下来的系列文章中，您将看到如何使用 Kubernetes 和 Vertex AI 扩展这些部署。
