Help us understand the problem. What is going on with this article?

TensorRTを用いたインファレンスの高速化

More than 1 year has passed since last update.

自己紹介

当記事はKichigai-Friends Advent Calendar 2017です.  

こんにちは,TrsNium です.
現在 大学を休学して放浪しています. 養ってくれる方を募集しています.

昨日開催されていたGTCJapnaで行って気になった,TensorRTの紹介をしたいと思います.
もし間違えがあればご指摘をお願いしいます.

TensorRTとは

Image
TensorRTはTensorFlowやPyTorchを用いいて学習したモデルを最適化をし,高速にインファレンスをすることを可能にすることができます.結果的にリアルタイムで動くアプリケーションに組み込むことでスループットの向上を狙うことができます.

Image

実際にTesla V100とTensorFlowを使った場合,秒間305枚をインファレンスできるのに対して,Tesla V100とTensorRTは秒間5707枚をインファレンスすることができます.何より1世代前のGPUが最新のGPUを使った場合よりもスループットが高いことに驚きです.
こんなもの使わないわけにはいかないでしょう…

何ができるのか

・全結合や畳み込みなどのレイヤーを簡潔なものに置き換えて最適化をすることができる.
・実行時間によって最適な並列アルゴリズムとレイヤーを選択することができる.
・FP32で学習したモデルをFP16やINT8を使い,精度を低下を最小限に抑えながら最適化をすることができる.
・並列的にデータを処理することができ,複数の入力を受けれるようすることができる.
・1つのカーネル実行のために連続するノードを単一のノードに結合することによって,GPUの使用率を改善し,メモリの記憶域と帯域幅を最適化することができる.

TensorFlowで学習したモデルをTensorRTで最適化をしてみる

TensorRTの導入

TensorRTの導入ですが,環境によって差があるので公式ドキュメンを見ていきましょう.
TensorRT3を使用しますが,その際に以下のものを必要とするので入れておきましょう.
・CUDA Toolkit 8.0 または 9.0.
・cuDNN 7.0. 以上
・python 2. または 3.

またTensorflowはver1.4を使用しています

実際のコード

コードの全体はこちらに置いておきます.
- 手順
- TensorFlowで学習したモデルをUFFフォーマットに変換
- 読み込んだモデルを元にTensorRT Engineを作る
- Engineを使いインファレンス
- 次回から簡単に読み込めるようにTensorRT Engineの保存と読み込み

1- TensorFlowでモデルを見込みUFFフォーマットに変換をする

import tensorflow as tf
import tensorrt as trt
from tensorrt.parsers import uffparser

import pycuda.driver as cuda
import pycuda.autoinit

#出力のnodeの名前
OUT_NAMES = ["inference/softmax"]

with tf.Session(config=config) as sess:
    saver = tf.train.Saver(tf.global_variables())
    saver.restore(sess, 'save/model.ckpt')

    graph_def = sess.graph_def()
    frozen_graph = tf.graph_util.convert_variables_to_constants(sess, graph_def, OUT_NAMES)
    tf_model = tf.graph_util.remove_training_nodes(frozen_graph)

#UFF modelへの変換
uff_model = uff.from_tensorflow(tf_model, OUT_NAMES)

2- 読み込んだモデルを元にTensorRT Engineを作る

# TensorRT EngineのためのUFF Streamを作る
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)

# uff parserを作り,モデルの入出力に関する情報を加える
parser = uffparser.create_uff_parser()
# (channel, im_size, im_size)
parser.register_input("Placeholder", (1,28,28), 0)
parser.register_output("inference/softmax")

# utility関数を用いてエンジンを作る(最後の引数はmax batch size と max workspace size)
engine = trt.utils.uff_to_trt_engine(G_LOGGER, uff_model, parser, MAX_BATCH_SIZE, MAX_WORKSPACE_SIZE)
parser.destroy()

3- Engineを使いインファレンス

runtime = trt.infer.create_infer_runtime(G_LOGGER)
context = engine.create_execution_context()

#データ読み込み(1件のみ)
img, label = self.data_set.test.next_batch(1)
img = img[0].astype(np.float32)
label = label[0]

# GPUにメモリ割り当てと,CPUにメモリ割り当て(推測後の結果を扱うために)
# 結果受取用の変数
output = np.empty(10, dtype = np.float32)

#alocate device memory (The size of the allocations is the size of the input and expected output * the batch size.)
d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)

bindings = [int(d_input), int(d_output)]

#推測をするためのcuda Streamを作成
stream = cuda.Stream()

#データをGPUに,推測と結果のコピー
cuda.memcpy_htod_async(d_input, img, stream)
#推測
context.enqueue(1, bindings, stream.handle, None)
#結果のコピー
cuda.memcpy_dtoh_async(output, d_output, stream)
#スレッドの同期
stream.synchronize()

-4 次回から簡単に読み込めるようにTensorRT Engineの保存と読み込み
次回からTFモデルからUFFモデルへの変換などしなくてもよくなります.

#engine の保存
trt.utils.write_engine_to_file("./tf_mnist.engine", engine.serialize())

# engine の読み込み
engine = trt.utils.load_engine(G_LOGGER, "./tf_mnist.engine")

まだ途中なので、後で訂正と追記をしていく予定です.
新しい情報などがあれば教えていただけると幸いです.

参考文献

TensorRT (https://developer.nvidia.com/tensorrt)
TensorRT Python API Documents (http://docs.nvidia.com/deeplearning/sdk/tensorrt-api/topics/topics/index.html)
Cuda9+cuDNN7+TensorFlowの導入(https://qiita.com/kumonkumon/items/3260ec4b4fa5194b1c9f)

Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
No comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
ユーザーは見つかりませんでした