TensorRT (3) TF-TRT (TensorFlow integaration with TensorRT) 소개 및 구현 예제

이전 글

2022.06.14 - [Deep Learning/Optimization (Algorithm)] - TensorRT (1) 개념, 최적화 방법, Workflow (Layer Fusion, Quantization 등)

2022.07.08 - [Deep Learning/Optimization (Algorithm)] - TensorRT (2) 설치 및 샘플 테스트 (Ubuntu 18.04 기준)

 

이전 글들에서 TensorRT의 기본 개념 및 설치 방법에 대하여 알아보았다. 이번 포스터에서는 TensorFlow 딥러닝 네트워크를 TensorRT로 최적화하는 방법에 대하여 정리하고자 한다. TensorFlow의 TensorRT integration 문서[1]를 참고하여 정리하고자 한다.

 

 

 

TF-TRT

TensorFlow™ integration with TensorRT™ (TF-TRT)는 TensorFlow(TF) 네트워크의 TensorRT(TRT)를 이용한 최적화를 위해 제공되는 module이다. TensorRT는 제한된 layer(e.x, convolution, fully connected, etc.)에 대한 최적화만을 진행하기에, 모든 TensorFlow 모델을 완벽히 최적화할 수는 없다. 따라서 TF-TRT는 TRT가 최적화할 수 있는 layer들에 대해 TRT 최적화를 하고, 나머지 layer들은 TF가 실행하도록 한다. 구체적으로 TRT는 TF 모델을 parsing하여 TF graph(=네트워크)의 가능한 모든 layer에 최적화를 진행한다.

 

TF-TRT를 통한 TF 모델 최적화의 간략하 workflow는 다음과 같다. TensorFlow 프레임워크를 이용해 학습된 네트워크를 saved_model 형식으로 저장한다. 이것을 TF-TRT의 python API를 통해 최적화한 후 또 다른 saved_model로 저장한다. Optimized saved_model을 이용해 TF에서 inference를 수행한다. Fig 1.은 TF와 TF-TRT의 workflow를 비교한 것이다. 한 단계를 추가하는 것만으로도 TF-TRT를 통한 inference 최적화가 되는 것을 확인할 수 있다.

 

Fig 1. TF와 TF-TRT의 workflow 비교 [1]

 

TF-TRT 사용법

TensorFlow 2 버전부터는 다음 코드로 TF pb 파일을 TRT로 최적화할 수 있다.

 

### Import convertet. TF-TRT는 TF의 binary에 포함되어 배포된다.
from tensorflow.python.compiler.tensorrt import trt_convert as trt

### converter class 초기화. input_saved_model_dir는 .pb 파일이 있는 경로
converter = trt.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir)

### 최적화 실행
converter.convert()

### 최적화된 네트워크 저장 (.pb 형식으로)
converter.save(output_saved_model_dir)

 

 

 

Sample Code

https://github.com/jhson989/tf-to-trt

GitHub - jhson989/tf-to-trt

위 github repository에 TF-TRT를 사용한 예제를 올려두었다. Sign language MNIST 데이터셋을 통해 학습한 임의 형태 CNN 네트워크를 TF-TRT를 통해 가속하는 예제이다.

 

TF-TRT의 가장 기본 configuration만을 사용해 만든 예제이다. 앞으로 build 개념 및 quantization, mixed precision 등 추가적인 TF-TRT 최적화를 적용해 볼 계획이다.

 

 

 

Reference

[1] https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html