Sirviendo modelos de TensorFlow con operaciones personalizadas

TensorFlow viene prediseñado con una extensa biblioteca de operaciones y núcleos de operaciones (implementaciones) optimizados para diferentes tipos de hardware (CPU, GPU, etc.). Estas operaciones se vinculan automáticamente al binario TensorFlow Serving ModelServer sin que el usuario requiera trabajo adicional. Sin embargo, hay dos casos de uso que requieren que el usuario vincule las operaciones al ModelServer explícitamente:

  • Ha escrito su propia operación personalizada (por ejemplo, usando esta guía )
  • Estás utilizando una operación ya implementada que no se incluye con TensorFlow.

Independientemente de si implementó la operación o no, para ofrecer un modelo con operaciones personalizadas, necesita acceso a la fuente de la operación. Esta guía lo guía a través de los pasos para usar la fuente para que las operaciones personalizadas estén disponibles para su publicación. Para obtener orientación sobre la implementación de operaciones personalizadas, consulte el repositorio de tensorflow/custom-op .

Requisito previo: con Docker instalado, ha clonado el repositorio de TensorFlow Serving y su directorio de trabajo actual es la raíz del repositorio.

Copie sobre el código fuente en el proyecto de servicio

Para construir TensorFlow Serving con sus operaciones personalizadas, primero deberá copiar la fuente de operaciones en su proyecto de servicio. Para este ejemplo, utilizará tensorflow_zero_out del repositorio de operaciones personalizadas mencionado anteriormente.

Dentro del repositorio de servicio, cree un directorio custom_ops , que albergará todas sus operaciones personalizadas. Para este ejemplo, solo tendrás el código tensorflow_zero_out .

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

Construya una biblioteca estática para la operación

En el archivo BUILD de tensorflow_zero_out, ve un objetivo que produce un archivo de objeto compartido ( .so ), que cargaría en Python para crear y entrenar su modelo. TensorFlow Serving, sin embargo, vincula estáticamente las operaciones en el momento de la compilación y requiere un archivo .a . Por lo tanto, agregará una regla de compilación que produzca este archivo en 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",
    ]
)

Construya ModelServer con la operación vinculada en

Para servir un modelo que utiliza una operación personalizada, debe compilar el binario de ModelServer con esa operación vinculada. Específicamente, agrega el objetivo de compilación zero_out_ops creado anteriormente al archivo BUILD de ModelServer.

Edite tensorflow_serving/model_servers/BUILD para agregar su objetivo de compilación de operaciones personalizado a SUPPORTED_TENSORFLOW_OPS que está incluido en el objetivo server_lib :

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

Luego use el entorno Docker para construir ModelServer:

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

Ofrecer un modelo que contenga su operación personalizada.

Ahora puede ejecutar el binario ModelServer y comenzar a servir un modelo que contenga esta operación personalizada:

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>

Enviar una solicitud de inferencia para probar la operación manualmente

Ahora puede enviar una solicitud de inferencia al servidor modelo para probar su operación personalizada:

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

Esta página contiene una API más completa para enviar solicitudes REST al servidor modelo.