在 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 擴充套件這些部署。
