TensorRTを使ってcaffeモデルを最適化するの巻

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()