TensorFlow r1.7にて、TensorRT統合

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

Google Develoers Blog、Announcing TensorRT integration with TensorFlow 1.7によると、
先日リリースされた TensorFlow r1.7 にて、NVIDIAのTensorRT が統合されました。

NVIDIAの協力で、TensorRT 3.0.4 と TensorFlow r1.7 が統合されました。

TensorFlowに実装した「ResNet-50」を CPUのFP32で実行したところ、11枚/秒でしたが、
NVIDIAのV100 FP32では、325枚/秒、V100 Tensor Cores With TensorRTでは、2657枚/秒だったようです。

TensorRTは、推論用のツールなので、TensorFlowで学習後、そのグラフを freeze します。
その後に、TensorRT用に最適化してからTensorRTを使って推論するというフローになります。

引用
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = number_between_0_and_1)
    ...  
    trt_graph = trt.create_inference_graph(
                 input_graph_def = frozen_graph_def, 
                 outputs = output_node_name,
                 max_batch_size=batch_size,
                 max_workspace_size_bytes=workspace_size,
                 precision_mode=precision)  # Get optimized graph

上記のように、frozenグラフを trt.create_inference_graph で TensorRT用の推論グラフに変換します。
この時、Max Batch Sizeも指定します。
まだ、TensorRTで使用するメモリサイズ、モデルで使用するサイズ(FP32, FP16, INT8)を指定します。

pool1 => block1 => block2 => block3 => block4 => pool5 => logits までを my_trt_op0 に統合し、最適化しています。

ちなみに、サポートしている TensorRT のバージョンは、3.0.4 です。