소개

Keras Applications의 모델을 이용하여 TensorRT용 모델로의 변환·추론까지를 대략 쓰고 있습니다.
TensorRT 등의 설치는 NVIDIA GPU Cloud의 컨테이너를 이용하므로 스코프외입니다.

준비

사전 조건

Docker 설치됨

GPU 컨테이너를 사용할 수 있습니다

실행 환경 시작

아래에서 실행 환경을 시작합니다.
Jupyter가 설치되었으므로 이후 코드는 Jupyter에서 시도 할 수 있습니다.

docker run -it --rm --gpus all -p 8888:8888 --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/tensorflow:20.02-tf2-py3 bash

# (Optional)
jupyter lab

nvcr.io/nvidia/tensorflow 는 NVIDIA GPU Cloud 에 등록되어 있는 컨테이너입니다.

이 컨테이너에는 TensorFlow 2.1, TensorRT 7.0, 기타 Jupyter 등이 설치되어 있습니다.

(Optional) shm-size, ulimit : 모델 변환시 메인 메모리를 엄청나게 사용하므로 메모리 확보 실패 대책도 겸해 설정한다.

(Optional) 8888 포트는 Jupyter 용입니다.

참고: Growth 설정

Keras를 사용할 때 자주 나오는 에러 대책의 메모리 제약.

메모리 제한 예

import tensorflow as tf

for dev in tf.config.experimental.list_physical_devices('GPU'):
    tf.config.experimental.set_memory_growth(dev, True)

변환

먼저 대상 모델을 출력합니다. 포인트는 저장 형식입니다.save_format 에 tf 를 지정하고 Tensorflow SavedModel에 저장합니다.
(TensorFlow 2.X에서는 디폴트이므로 불필요합니다만)

Keras에서 모델 저장

from tensorflow.keras.applications.vgg16 import VGG16

model = VGG16(weights='imagenet')
model.save('./vgg16', save_format='tf')

그런 다음 모델을 변환합니다.

TensorRT 변환 단정밀도 부동 소수점(Float32)판

from tensorflow.python.compiler.tensorrt import trt_convert as trt

converter = trt.TrtGraphConverterV2(input_saved_model_dir='./vgg16',
                                    conversion_params=trt.DEFAULT_TRT_CONVERSION_PARAMS)
converter.convert()
converter.save('./vgg16-tensorrt')

반 정밀도 부동 소수점 버전

Float16으로 변환하려면 변환기의 매개 변수를 변경합니다.

converter = trt.TrtGraphConverterV2(input_saved_model_dir='./vgg16',
                                    conversion_params=trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(precision_mode=trt.TrtPrecisionMode.FP16))

정수판

8비트 정수로 설정하면 보정이 필요합니다. 학습에 사용한 데이터를 이용하면 좋다고 생각합니다.

이번 VGG16의 설정에서는 $(N, 224, 224, 3)$의 Shape로 전달합니다.

import numpy as np

def calibration_input_fn():  # キャリブレーション用データ生成関数
    yield np.random.uniform(size=(5, 224, 224, 3)).astype(np.float32),  # 末尾の , を忘れないこと

converter = trt.TrtGraphConverterV2(input_saved_model_dir='./vgg16',
                                    conversion_params=trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(precision_mode=trt.TrtPrecisionMode.INT8, use_calibration=True))
converter.convert(calibration_input_fn=calibration_input_fn)
converter.save('./vgg16-tensorrt')

추론

변환한 모델을 로드해, 추론에 이용하는 오브젝트를 꺼냅니다.
그리고 그 객체를 함수를 호출하면 추론이 실행됩니다.

model = tf.saved_model.load('./vgg16-tensorrt', tags=[tf.saved_model.SERVING])
infer = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

# ダミーの入力
x = np.random.uniform(size=(3, 224, 224, 3)).astype(np.float32)
# 推論
y = infer(tf.convert_to_tensor(x))['predictions']

참고 : 입력 Shape는 다음과 같이 얻을 수 있습니다.

infer.inputs[0].shape
>>> TensorShape([None, 224, 224, 3])

실행 결과 비교

마지막으로 대략적입니다만, 실행 결과의 비교(위가 TensorRT)가 됩니다. 이런 적당하더라도 실행 속도가 향상됩니다.
메모리 사용량도 감소하기 때문에, 복수 모델 실행 가능 등 여러가지 할 수 있을 것 같고, 실행 환경의 GPU 아키텍쳐에서도 혜택의 정도가 크게 바뀐다고 생각합니다.

실행 환경: GeForce GTX 1080 Ti