使用TF Serving在Hugging Face中部署TensorFlow视觉模型
在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的serving签名：
saved_model_cli show --dir {temp_model_dir}/saved_model/1 --tag_set serve --signature_def serving_default这应该输出：
给定的SavedModel SignatureDef包含以下输入：
inputs['pixel_values'] tensor_info：
dtype：DT_FLOAT
shape：(-1, -1, -1, -1)
name：serving_default_pixel_values:0
给定的SavedModel SignatureDef包含以下输出：
outputs['logits'] tensor_info：
dtype：DT_FLOAT
shape：(-1, 1000)
name：StatefulPartitionedCall:0
方法名称是：tensorflow/serving/predict可以注意到，模型接受单个4维输入（即pixel_values），其轴如下：(batch_size, num_channels, height, width)。对于此模型，可接受的高度和宽度设置为224，通道数为3。可以通过检查模型的config参数（model.config）来验证这一点。模型生成一个1000维的logits向量。
通常，每个ML模型都有特定的预处理和后处理步骤。ViT模型也不例外。主要的预处理步骤包括：
-
将图像像素值缩放到[0, 1]范围内。
-
将缩放后的像素值归一化为[-1, 1]。
-
调整图像大小，使其具有(224, 224)的空间分辨率。
您可以通过调查与模型关联的特征提取器来确认这些：
from transformers import AutoFeatureExtractor
feature_extractor = AutoFeatureExtractor.from_pretrained(ckpt)
print(feature_extractor)这应该打印：
ViTFeatureExtractor {
"do_normalize": true,
"do_resize": true,
"feature_extractor_type": "ViTFeatureExtractor",
"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=feature_extractor.image_mean, std=feature_extractor.image_std
):
# 首先缩放到 [0, 1] 的值范围，然后进行归一化。
img = img / 255
mean = tf.constant(mean)
std = tf.constant(std)
return (img - mean) / std您还需要调整图像的大小并转置，以便具有前导通道维度，以遵循 🤗 Transformers 的标准格式。下面的代码片段显示了所有的预处理步骤：
CONCRETE_INPUT = "pixel_values" # 这是我们通过 SavedModel CLI 调查的内容。
SIZE = feature_extractor.size
def normalize_img(
img, mean=feature_extractor.image_mean, std=feature_extractor.image_std
):
# 首先缩放到 [0, 1] 的值范围，然后进行归一化。
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)
) # 因为 HF 模型是 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
给定的SavedModel SignatureDef包含以下输入：
inputs['string_input'] tensor_info：
dtype: DT_STRING
shape: (-1)
name: serving_default_string_input:0
给定的SavedModel SignatureDef包含以下输出：
outputs['confidence'] tensor_info：
dtype: DT_FLOAT
shape: (-1)
name: StatefulPartitionedCall:0
outputs['label'] tensor_info：
dtype: DT_STRING
shape: (-1)
name: StatefulPartitionedCall:1
方法名是：tensorflow/serving/predict你可以注意到模型的签名已经改变。具体来说，输入类型现在是字符串，模型返回两个结果：置信度和字符串标签。
假设你已经安装了TF Serving（在Colab笔记本中有介绍），现在你可以部署这个模型了！
只需要执行一个命令：
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）的模型。
回想一下，你导出的模型可以接受使用base64格式编码的字符串输入。因此，你可以按照以下方式构建请求载荷：
# 获取一张可爱猫咪的图片。
image_path = tf.keras.utils.get_file(
"image.jpg", "http://images.cocodataset.org/val2017/000000039769.jpg"
)
# 从磁盘中读取原始字节并进行编码。
bytes_inputs = tf.io.read_file(image_path)
b64str = base64.urlsafe_b64encode(bytes_inputs.numpy()).decode("utf-8")
# 创建请求载荷。
data = json.dumps({"signature_name": "serving_default", "instances": [b64str]})TF Serving的REST端点的请求载荷格式规范可以在这里找到。你可以在instances中传递多个编码图像。这种端点适用于在线预测场景。对于包含多个数据点的输入，你可以启用批处理以获得性能优化的好处。
现在你可以调用API：
headers = {"content-type": "application/json"}
json_response = requests.post(
"http://localhost:8501/v1/models/vit:predict", data=data, headers=headers
)
print(json.loads(json_response.text))
# {'predictions': [{'label': 'Egyptian cat', 'confidence': 0.896659195}]}REST API的地址是 – http://localhost:8501/v1/models/vit:predict，遵循这里的规范。默认情况下，它总是选择最新版本的模型。但如果你想要特定版本，可以使用这样的地址： http://localhost:8501/v1/models/vit/versions/1:predict。
虽然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 秒超时
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 扩展这些部署。
-
gRPC
-
计算机视觉的实用机器学习
-
在 Hugging Face Transformers 中加速 TensorFlow 模型
