TF Servingを使用してHugging FaceでTensorFlow Visionモデルを展開する

Hugging FaceでTensorFlow VisionモデルをTF Servingを使って展開する

過去数ヶ月間、Hugging Faceチームと外部の貢献者は、TransformersにさまざまなビジョンモデルをTensorFlowで追加しました。このリストは包括的に拡大しており、ビジョントランスフォーマー、マスク付きオートエンコーダー、RegNet、ConvNeXtなど、最先端の事前学習モデルがすでに含まれています！

TensorFlowモデルを展開する際には、さまざまな選択肢があります。使用ケースに応じて、モデルをエンドポイントとして公開するか、アプリケーション自体にパッケージ化するかを選択できます。TensorFlowには、これらの異なるシナリオに対応するツールが用意されています。

この投稿では、TensorFlow Serving（TF Serving）を使用してローカルでビジョントランスフォーマーモデル（画像分類用）を展開する方法を紹介します。これにより、開発者はモデルをRESTエンドポイントまたはgRPCエンドポイントとして公開できます。さらに、TF Servingはモデルのウォームアップ、サーバーサイドバッチ処理など、多くの展開固有の機能を提供しています。

• プライベートハブのご紹介：機械学習を活用した新しいビルド方法
• ディープダイブ：Hugging Face Optimum GraphcoreにおけるビジョンTransformer
• Megatron-LMを使用して言語モデルをトレーニングする方法

この投稿全体で示される完全な動作するコードを取得するには、冒頭に示されているColabノートブックを参照してください。

🤗 TransformersのすべてのTensorFlowモデルには、save_pretrained()というメソッドがあります。このメソッドを使用すると、モデルの重みをh5形式およびスタンドアロンのSavedModel形式でシリアライズできます。TF Servingでは、モデルをSavedModel形式で提供する必要があります。そこで、まずビジョントランスフォーマーモデルをロードして保存します。

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

これは次のように出力されます：

与えられた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）を受け入れることがわかります。これには次の軸があります：(バッチサイズ、チャネル数、高さ、幅)。このモデルでは、許容される高さと幅は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 としてシリアライズする必要があります。

前処理

前処理では、画像の正規化が最も重要なコンポーネントの1つです：

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

また、画像をリサイズし、チャンネルの次元が先行するようにトランスポーズする必要があります。以下のコードスニペットには、すべての前処理ステップが表示されています：

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 モデルはチャンネルが先行です。
    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

まず、モデルの forward pass メソッド（ call() ）から具体的な関数を派生させることで、モデルをグラフにきれいにコンパイルします。その後、以下の手順を順番に適用できます：

1. 入力を前処理操作に通します。

2. 前処理された入力を派生した具体的な関数に通します。

3. 出力を後処理し、見やすい形式の辞書で返します。

さあ、モデルをエクスポートしましょう！

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

モデルのシグネチャが変更されたことに気付くことができます。具体的には、入力のタイプが文字列になり、モデルは2つの要素を返します：信頼度スコアと文字列のラベル。

TF Servingをすでにインストールしている場合（Colabノートブックで説明されています）、このモデルを展開する準備が整いました！

これを行うには、たった1つのコマンドが必要です：

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はgRPCエンドポイントに対して8500ポートを使用します。

• model_nameは、APIを呼び出すために使用されるモデル名を指定します（何でも構いません）。

• model_base_pathは、TF Servingがモデルの最新バージョンをロードするために使用するベースモデルパスを示します。

（サポートされているパラメータの完全なリストはこちらです。）

そして、できあがり！数分以内に、RESTエンドポイントとgRPCエンドポイントを持つ展開済みのモデルで実行できるようになります。

モデルは、ベース64形式でエンコードされた文字列入力を受け入れるようにエクスポートされていることを思い出してください。したがって、リクエストペイロードを作成するために次のようなことができます：

# 可愛い猫の画像を取得します。
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: "エジプトの猫"
  }
}
model_spec {
  name: "resnet"
  version {
    value: 1
  }
  signature_name: "serving_default"
}

また、上記の結果から興味のあるキーと値を取得することもできます：

grpc_predictions.outputs["label"].string_val, grpc_predictions.outputs[
    "confidence"
].float_val
# ([b'エジプトの猫'], [0.8966591954231262])

この投稿では、TF Servingを使用してTransformersからTensorFlowビジョンモデルをデプロイする方法を学びました。ローカルデプロイは週末のプロジェクトには最適ですが、これらのデプロイを多くのユーザーに提供するためにスケーリングする必要があります。次のシリーズの投稿では、KubernetesとVertex AIを使用してこれらのデプロイをスケールアップする方法を紹介します。

• gRPC

• コンピュータビジョンのための実用的な機械学習

• Hugging Face Transformersにおける高速なTensorFlowモデル

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!