Servindo modelos do TensorFlow com operações personalizadas

O TensorFlow vem pré-construído com uma extensa biblioteca de operações e kernels operacionais (implementações) ajustados para diferentes tipos de hardware (CPU, GPU, etc.). Essas operações são automaticamente vinculadas ao binário TensorFlow Serving ModelServer sem nenhum trabalho adicional exigido pelo usuário. No entanto, há dois casos de uso que exigem que o usuário vincule explicitamente as operações ao ModelServer:

  • Você escreveu sua própria operação personalizada (por exemplo, usando este guia )
  • Você está usando uma operação já implementada que não vem com o TensorFlow

Independentemente de você ter implementado a operação ou não, para servir um modelo com operações personalizadas, você precisa de acesso à origem da operação. Este guia orienta você nas etapas de uso da fonte para disponibilizar operações personalizadas para veiculação. Para obter orientação sobre a implementação de operações personalizadas, consulte o repositório tensorflow/custom-op .

Pré-requisito: com o Docker instalado, você clonou o repositório do TensorFlow Serving e seu diretório de trabalho atual é a raiz do repositório.

Copie a fonte de operação para o projeto Serving

Para criar o TensorFlow Serving com suas operações personalizadas, primeiro você precisa copiar a origem da operação para seu projeto de serviço. Para este exemplo, você usará tensorflow_zero_out do repositório de operações personalizadas mencionado acima.

Dentro do repositório de serviço, crie um diretório custom_ops , que abrigará todas as suas operações personalizadas. Para este exemplo, você terá apenas o código tensorflow_zero_out .

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

Crie uma biblioteca estática para a operação

No arquivo BUILD do tensorflow_zero_out, você vê um alvo produzindo um arquivo de objeto compartilhado ( .so ), que você carregaria em python para criar e treinar seu modelo. O TensorFlow Serving, no entanto, vincula operações estaticamente no momento da compilação e requer um arquivo .a . Então você adicionará uma regra de construção que produz este arquivo para 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",
    ]
)

Construa ModelServer com a operação vinculada

Para servir um modelo que usa uma operação personalizada, você deve construir o binário ModelServer com essa operação vinculada. Especificamente, você adiciona o destino de construção zero_out_ops criado acima ao arquivo BUILD do ModelServer.

Edite tensorflow_serving/model_servers/BUILD para adicionar seu destino de compilação de operação personalizado a SUPPORTED_TENSORFLOW_OPS que está incluído no destino server_lib :

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

Em seguida, use o ambiente Docker para construir o ModelServer:

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

Sirva um modelo contendo sua operação personalizada

Agora você pode executar o binário ModelServer e começar a servir um modelo que contém esta operação 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>

Envie uma solicitação de inferência para testar a operação manualmente

Agora você pode enviar uma solicitação de inferência ao servidor modelo para testar sua operação personalizada:

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

Esta página contém uma API mais completa para enviar solicitações REST ao servidor modelo.