Pythonにおいて、caffeモデルを最適化するためにTensorRTを使用しているこのチュートリアルをやってみた。

TensorRTを使ってcaffeモデルを最適化¶

TensorRT 4.0は、実行・保存可能なカフェモデルを読み込んで最適化するためのPython APIをサポートしている。先ず、必要なモジュールをインポートする。

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from random import randint
from PIL import Image
from matplotlib.pyplot import imshow #to show test case
from tensorrt import parsers

通常、最初にやる事は、モデル変換と推論プロセスの間に頻繁に使われるロガーを作成することで、tensorrt.infer.ConsoleLoggerで簡単なロガー実装を提供しているが、自分でもロガーを定義することができる。

G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)

次に、MNIST datasetの数字を分類するのに使うモデル用の定数を定義する。

INPUT_LAYERS = ['data']
OUTPUT_LAYERS = ['prob']
INPUT_H = 28
INPUT_W =  28
OUTPUT_SIZE = 10

次に、モデルデータ用のデータパスを設定する。

MODEL_PROTOTXT = '/home/workspace/TensorRT-4.0.1.6/data/mnist/mnist.prototxt'
CAFFE_MODEL = '/home/workspace/TensorRT-4.0.1.6/data/mnist/mnist.caffemodel'
DATA = '/home/workspace/TensorRT-4.0.1.6/data/mnist/'
IMAGE_MEAN = '/home/workspace/TensorRT-4.0.1.6/data/mnist/mnist_mean.binaryproto'

最初の工程でエンジンを作成する。Python APIが、この工程を簡素化するユーティリティを提供してくれる。ここで、tensorrt.utilsのカフェモデル変換ユーティリティを使用する。ロガー、model prototxtへのパス、モデルファイル、最大バッチサイズ、最大ワークスペースサイズ、アウトプットレイヤー、ウェイトのデータタイプを提供する。

engine = trt.utils.caffe_to_trt_engine(G_LOGGER,
                                       MODEL_PROTOTXT,
                                       CAFFE_MODEL,
                                       1,
                                       1 << 20,
                                       OUTPUT_LAYERS,
                                       trt.infer.DataType.FLOAT)

上で作成したエンジン用のテストケースを生成する。

rand_file = randint(0,9)
path = DATA + str(rand_file) + '.pgm'
im = Image.open(path)
%matplotlib inline
imshow(np.asarray(im))
arr = np.array(im)
img = arr.ravel()
print("Test Case: " + str(rand_file))

Test Case: 7

次に、インプット画像にmeanを適用、caffeparserが読み込みできる.binaryprotoファイルにこれは保存されている。

parser = parsers.caffeparser.create_caffe_parser()
mean_blob = parser.parse_binary_proto(IMAGE_MEAN)
parser.destroy()
#NOTE: This is different than the C++ API, you must provide the size of the data
mean = mean_blob.get_data(INPUT_W ** 2)
data = np.empty([INPUT_W ** 2])
for i in range(INPUT_W ** 2):
    data[i] = float(img[i]) - mean[i]
mean_blob.destroy()

その次に、推論用ランタイムとエンジン用の実行コンテクストを作成する。

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

これで推論を実行できるようになり、先ず最初に、このモデルのデータが正確なデータタイプ(FP32)であることを確かめる。それから、CPUに推論結果を保存するための空のアレイを割り当てる。

assert(engine.get_nb_bindings() == 2)
#convert input data to Float32
img = img.astype(np.float32)
#create output array to receive data
output = np.empty(OUTPUT_SIZE, dtype = np.float32)

次に、PyCUDAを使ってGPUメモリに割り当ててエンジンに登録する。これらの割り当てサイズは、インプット/予想アウトプット×バッチサイズになる。

d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)

エンジンは、バインディング(GPUメモリへのポインター)を要求する。PyCUDAはメモリ割り当ての結果をintにキャスティングすることでこの事を可能にしている。

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

推論実行のためにcudaストリームを作成する。

stream = cuda.Stream()

次に、データをGPUに転送して、推論を実行して。結果をCPUにコピーする。

#transfer input data to device
cuda.memcpy_htod_async(d_input, img, stream)
#execute model
context.enqueue(1, bindings, stream.handle, None)
#transfer predictions back
cuda.memcpy_dtoh_async(output, d_output, stream)
#syncronize threads
stream.synchronize()

推論結果を得るのにnp.argmaxが利用できる。

print("Test Case: " + str(rand_file))
print ("Prediction: " + str(np.argmax(output)))

Test Case: 7
Prediction: 7

エンジンは後で使えるように保存することも可能だ。

trt.utils.write_engine_to_file("/home/workspace/TensorRT-4.0.1.6/data/mnist/new_mnist.engine", engine.serialize())

True

tensorrt.utils.load_engineを使って後でエンジンをロードする事ができる。

new_engine = trt.utils.load_engine(G_LOGGER, "/home/workspace/TensorRT-4.0.1.6/data/mnist/new_mnist.engine")

最後に、コンテクスト、エンジン、ランタイムを破棄する。

context.destroy()
engine.destroy()
new_engine.destroy()
runtime.destroy()