Diffuser des modèles TensorFlow avec des opérations personnalisées

TensorFlow est livré pré-construit avec une vaste bibliothèque d'opérations et de noyaux opérationnels (implémentations) optimisés pour différents types de matériel (CPU, GPU, etc.). Ces opérations sont automatiquement liées au binaire TensorFlow Serving ModelServer sans aucun travail supplémentaire requis par l'utilisateur. Cependant, il existe deux cas d'utilisation qui nécessitent que l'utilisateur se connecte explicitement aux opérations dans le ModelServer :

  • Vous avez écrit votre propre opération personnalisée (ex. en utilisant ce guide )
  • Vous utilisez une opération déjà implémentée qui n'est pas fournie avec TensorFlow

Que vous ayez implémenté l'opération ou non, pour servir un modèle avec des opérations personnalisées, vous devez accéder à la source de l'opération. Ce guide vous guide à travers les étapes d'utilisation de la source pour rendre les opérations personnalisées disponibles pour la diffusion. Pour obtenir des conseils sur la mise en œuvre des opérations personnalisées, veuillez vous référer au dépôt tensorflow/custom-op .

Condition préalable : une fois Docker installé, vous avez cloné le référentiel TensorFlow Serving et votre répertoire de travail actuel est la racine du référentiel.

Copier la source opérationnelle dans le projet de service

Afin de créer TensorFlow Serving avec vos opérations personnalisées, vous devez d'abord copier la source de l'opération dans votre projet de diffusion. Pour cet exemple, vous utiliserez tensorflow_zero_out du référentiel d'opérations personnalisées mentionné ci-dessus.

Dans le dépôt de service, créez un répertoire custom_ops , qui hébergera toutes vos opérations personnalisées. Pour cet exemple, vous n'aurez que le code tensorflow_zero_out .

mkdir tensorflow_serving/custom_ops
cp -r <custom_ops_repo_root>/tensorflow_zero_out tensorflow_serving/custom_ops

Construire une bibliothèque statique pour l'opération

Dans le fichier BUILD de tensorflow_zero_out, vous voyez une cible produisant un fichier objet partagé ( .so ), que vous chargeriez dans python afin de créer et d'entraîner votre modèle. TensorFlow Serving, cependant, lie statiquement les opérations au moment de la construction et nécessite un fichier .a . Vous allez donc ajouter une règle de build qui produit ce fichier à tensorflow_serving/custom_ops/tensorflow_zero_out/BUILD :

cc_library(
    name = 'zero_out_ops',
    srcs = [
        "cc/kernels/zero_out_kernels.cc",
        "cc/ops/zero_out_ops.cc",
    ],
    alwayslink = 1,
    deps = [
        "@org_tensorflow//tensorflow/core:framework",
    ]
)

Construisez ModelServer avec l'opération liée

Pour servir un modèle qui utilise une opération personnalisée, vous devez créer le binaire ModelServer avec cette opération liée. Plus précisément, vous ajoutez la cible de build zero_out_ops créée ci-dessus au fichier BUILD du ModelServer.

Modifiez tensorflow_serving/model_servers/BUILD pour ajouter votre cible de build d'opération personnalisée à SUPPORTED_TENSORFLOW_OPS qui est incluse dans la cible server_lib :

SUPPORTED_TENSORFLOW_OPS = [
    ...
    "//tensorflow_serving/custom_ops/tensorflow_zero_out:zero_out_ops"
]

Utilisez ensuite l'environnement Docker pour créer le ModelServer :

tools/run_in_docker.sh bazel build tensorflow_serving/model_servers:tensorflow_model_server

Servir un modèle contenant votre opération personnalisée

Vous pouvez maintenant exécuter le binaire ModelServer et commencer à servir un modèle contenant cette opération personnalisée :

tools/run_in_docker.sh -o "-p 8501:8501" \
bazel-bin/tensorflow_serving/model_servers/tensorflow_model_server \
--rest_api_port=8501 --model_name=<model_name> --model_base_path=<model_base_path>

Envoyer une demande d'inférence pour tester l'opération manuellement

Vous pouvez maintenant envoyer une requête d'inférence au serveur de modèles pour tester votre opération personnalisée :

curl http://localhost:8501/v1/models/<model_name>:predict -X POST \
-d '{"inputs": [[1,2], [3,4]]}'

Cette page contient une API plus complète pour envoyer des requêtes REST au serveur de modèles.