Phục vụ các mô hình TensorFlow với các hoạt động tùy chỉnh

TensorFlow được xây dựng sẵn với một thư viện mở rộng gồm các op và hạt nhân op (triển khai) được tinh chỉnh cho các loại phần cứng khác nhau (CPU, GPU, v.v.). Các hoạt động này được tự động liên kết với tệp nhị phân TensorFlow Serve ModelServer mà người dùng không cần thực hiện thêm công việc nào. Tuy nhiên, có hai trường hợp sử dụng yêu cầu người dùng liên kết rõ ràng các hoạt động với ModelServer:

  • Bạn đã viết op tùy chỉnh của riêng mình (ví dụ: sử dụng hướng dẫn này )
  • Bạn đang sử dụng một op đã được triển khai nhưng không được cung cấp cùng với TensorFlow

Bất kể bạn có triển khai op hay không, để phục vụ một mô hình với các op tùy chỉnh, bạn cần có quyền truy cập vào nguồn của op. Hướng dẫn này sẽ hướng dẫn bạn các bước sử dụng nguồn để cung cấp các hoạt động tùy chỉnh để phân phát. Để được hướng dẫn về cách triển khai các hoạt động tùy chỉnh, vui lòng tham khảo kho lưu trữ tensorflow/custom-op .

Điều kiện tiên quyết: Sau khi cài đặt Docker, bạn đã sao chép kho lưu trữ TensorFlow Serve và thư mục làm việc hiện tại của bạn là thư mục gốc của kho lưu trữ.

Sao chép mã nguồn vào dự án Phục vụ

Để xây dựng Dịch vụ TensorFlow bằng các hoạt động tùy chỉnh của bạn, trước tiên bạn cần sao chép nguồn hoạt động vào dự án phân phối của mình. Trong ví dụ này, bạn sẽ sử dụng tensorflow_zero_out từ kho lưu trữ tùy chỉnh được đề cập ở trên.

Trong kho lưu trữ phục vụ, hãy tạo một thư mục custom_ops , thư mục này sẽ chứa tất cả các hoạt động tùy chỉnh của bạn. Trong ví dụ này, bạn sẽ chỉ có mã tensorflow_zero_out .

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

Xây dựng thư viện tĩnh cho op

Trong tệp BUILD của tensorflow_zero_out, bạn thấy mục tiêu tạo tệp đối tượng dùng chung ( .so ), tệp này bạn sẽ tải vào python để tạo và huấn luyện mô hình của mình. Tuy nhiên, TensorFlow Serve liên kết tĩnh các hoạt động tại thời điểm xây dựng và yêu cầu tệp .a . Vì vậy, bạn sẽ thêm quy tắc xây dựng tạo tệp này vào 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",
    ]
)

Xây dựng ModelServer với op được liên kết trong

Để phục vụ một mô hình sử dụng op tùy chỉnh, bạn phải xây dựng tệp nhị phân ModelServer với op được liên kết trong đó. Cụ thể, bạn thêm mục tiêu xây dựng zero_out_ops đã tạo ở trên vào tệp BUILD của ModelServer.

Chỉnh sửa tensorflow_serving/model_servers/BUILD để thêm mục tiêu xây dựng op tùy chỉnh của bạn vào SUPPORTED_TENSORFLOW_OPS được bao gồm trong mục tiêu server_lib :

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

Sau đó sử dụng môi trường Docker để xây dựng ModelServer:

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

Phục vụ một mô hình chứa op tùy chỉnh của bạn

Bây giờ bạn có thể chạy tệp nhị phân ModelServer và bắt đầu phân phối một mô hình có chứa op tùy chỉnh này:

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>

Gửi yêu cầu suy luận để kiểm tra op theo cách thủ công

Bây giờ bạn có thể gửi yêu cầu suy luận đến máy chủ mô hình để kiểm tra hoạt động tùy chỉnh của mình:

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

Trang này chứa API hoàn chỉnh hơn để gửi yêu cầu REST đến máy chủ mô hình.