Pythonにおいて、caffeモデルを最適化するためにTensorRTを使用しているこのチュートリアルをやってみた。
TensorRTを使ってcaffeモデルを最適化¶
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))
次に、インプット画像に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)))
エンジンは後で使えるように保存することも可能だ。
trt.utils.write_engine_to_file("/home/workspace/TensorRT-4.0.1.6/data/mnist/new_mnist.engine", engine.serialize())
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()