TensorFlow Serving 是用于機器學習模型的高性能靈活服務系統(tǒng)，而 NVIDIA TensorRT? 是實現(xiàn)高性能深度學習推理的平臺，通過將二者相結(jié)合，用戶便可獲得更高性能，從而輕松實現(xiàn) GPU 推理。TensorFlow 團隊與 NVIDIA 合作，在 TensorFlow v1.7 中首次添加了對 TensorRT 的支持。此后，我們一直密切合作，共同改進 TensorFlow-TensorRT 集成（稱為 TF-TRT）。TensorFlow Serving 1.13 現(xiàn)已實現(xiàn)這種集成，TensorFlow 2.0 很快也會進行集成。

在 之前的文章 中，我們介紹了如何借助 Docker 使用 TensorFlow Serving。在本文中，我們將展示以同樣的方式運行經(jīng) TF-TRT 轉(zhuǎn)換的模型有多簡單。與之前一樣，我們嘗試在生產(chǎn)環(huán)境中部署 ResNet 模型。下方所有示例均在配備 Titan-V GPU 的工作站上運行。

注：ResNet 鏈接

https://github.com/tensorflow/models/tree/master/official/resnet

在 GPU 上利用 TensorFlow Serving 部署 ResNet

在這項練習中，我們僅下載 經(jīng)過預訓練的 ResNet SavedModel：

$ mkdir /tmp/resnet

$ curl -s https://storage.googleapis.com/download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v2_fp32_savedmodel_NHWC_jpg.tar.gz | tar --strip-components=2 -C /tmp/resnet -xvz

$ ls /tmp/resnet

1538687457

注：經(jīng)過預訓練的 ResNet 鏈接

https://github.com/tensorflow/models/tree/master/official/resnet#pre-trained-model

在之前的文章中，我們展示了如何使用 TensorFlow Serving CPU Docker 圖像提供模型。在這里，我們運行 GPU Docker 圖像（請查看此處了解相關說明），以借助 GPU 提供并測試此模型：

$ docker pull tensorflow/serving:latest-gpu

$ docker run --rm --runtime=nvidia -p 8501:8501 --name tfserving_resnet \

-v /tmp/resnet:/models/resnet -e MODEL_NAME=resnet -t tensorflow/serving:latest-gpu &

…

… server.cc:286] Running gRPC ModelServer at 0.0.0.0:8500 …

… server.cc:302] Exporting HTTP/REST API at:localhost:8501 …

$ curl -o /tmp/resnet/resnet_client.py https://raw.githubusercontent.com/tensorflow/serving/master/tensorflow_serving/example/resnet_client.py

$ python /tmp/resnet/resnet_client.py

Prediction class:286, avg latency:18.0469 ms

注：此處 鏈接

https://www.tensorflow.org/serving/docker#serving_with_docker_using_your_gpu

此 docker run 命令會啟動 TensorFlow Serving 服務器，以提供 /tmp/resnet 中已下載的 SavedModel，并在主機中開放 REST API 端口 8501。resnet_client.py 會發(fā)送一些圖像給服務器，并返回服務器所作的預測。現(xiàn)在讓我們終止 TensorFlow Serving 容器的運行，以釋放所占用的 GPU 資源。

$ docker kill tfserving_resnet

注：REST API 鏈接

https://www.tensorflow.org/tfx/serving/api_rest

利用 TF-TRT 轉(zhuǎn)換和部署模型

現(xiàn)在，我們有了工作模型。為了享受 TensorRT 帶來的好處，我們需要在 TensorFlow Serving Docker 容器內(nèi)運行轉(zhuǎn)換命令，從而將此模型轉(zhuǎn)換為使用 TensorRT 運行運算的模型：

$ docker pull tensorflow/tensorflow:latest-gpu

$ docker run --rm --runtime=nvidia -it -v /tmp:/tmp tensorflow/tensorflow:latest-gpu /usr/local/bin/saved_model_cli \

convert --dir /tmp/resnet/1538687457 --output_dir /tmp/resnet_trt/1538687457 --tag_set serve \

tensorrt --precision_mode FP32 --max_batch_size 1 --is_dynamic_op True

在這里，我們運行了 saved_model_cli 命令行工具，其中內(nèi)置了對 TF-TRT 轉(zhuǎn)換的支持。--dir 和 --output_dir 參數(shù)會指示 SavedModel 的位置以及在何處輸出轉(zhuǎn)換后的 SavedModel，而 --tag_set 則指示 SavedModel 中要轉(zhuǎn)換的圖表。隨后，我們在命令行中傳遞 tensorrt 并指定配置，明確指示其運行 TF-TRT 轉(zhuǎn)換器：

--precision_mode 指示轉(zhuǎn)換器需使用的精度，目前其僅支持 FP32 和 FP16

--max_batch_size 指示輸入的批次大小上限。此轉(zhuǎn)換器要求將由 TensorRT 處理的所有張量將其首個維度作為批次維度，而該參數(shù)則指示推理過程中會產(chǎn)生的最大值。若已知推理過程中的實際批次大小上限且該值與之匹配，則轉(zhuǎn)換后的模型即為最優(yōu)模型。請注意，轉(zhuǎn)換后的模型無法處理批次規(guī)模大于此處所指定大小的輸入，但可處理批次規(guī)模更小的輸入

--is_dynamic_op 指示在模型運行時進行實際轉(zhuǎn)換。原因在于，在進行轉(zhuǎn)換時，TensorRT 需要明確所有形狀。對于本例中使用的 ResNet 模型，其張量沒有固定的形狀，因此我們需要此參數(shù)

注：saved_model_cli 鏈接

https://www.tensorflow.org/guide/saved_model#cli_to_inspect_and_execute_savedmodel

現(xiàn)在，我們只需為模型指定正確的目錄，便可利用 Docker 提供經(jīng) TF-TRT 轉(zhuǎn)換的模型，這與之前一樣簡單：

$ docker run --rm --runtime=nvidia -p 8501:8501 --name tfserving_resnet \

-v /tmp/resnet_trt:/models/resnet -e MODEL_NAME=resnet -t tensorflow/serving:latest-gpu &

…

… server.cc:286] Running gRPC ModelServer at 0.0.0.0:8500 …

… server.cc:302] Exporting HTTP/REST API at:localhost:8501 …

向其發(fā)送請求：

$ python /tmp/resnet/resnet_client.py

Prediction class:286, avg latency:15.0287 ms

最后，我們終止容器的運行：

$ docker kill tfserving_resnet

我們可以看到，使用 TensorFlow Serving 和 Docker 生成經(jīng) TF-TRT 轉(zhuǎn)換的模型與提供一般模型一樣簡單。此外，以上為展示內(nèi)容，其中的性能數(shù)字僅適用于我們所使用的模型和運行本示例的設備，但它的確展現(xiàn)出使用 TF-TRT 所帶來的性能優(yōu)勢。

TensorFlow 2.0 發(fā)布在即，TensorFlow 團隊和 NVIDIA 正在共同努力，以確保 TF-TRT 能在 2.0 中流暢運行。如需了解最新信息，請查看 TF-TRT GitHub 代碼庫(https://github.com/tensorflow/tensorrt)。