TENSOR RT - (2)

지난시간에 이어서, TENSOR RT RUNTIME API 에 대해 말씀드리겠습니다. 

저번 시간에서는 trt.Builder(Logger 설정) 을 통해, builder를 만들어 주었고, 이렇게 만들어진 builder를 통해, builder.create_network(common.BATCH_SIZE) 의 TRT.INETWORK 객체와 builder.create_builder_config()를 통해 TRT.CONFIG()를 생성해주었습니다. 그다음, 네트워크 생성 함수를 통해, Weight 와 INetwork를 일치시켜, network를 설정해주었습니다. 오늘은 그 이후의 단계에 대해 설명 드리겠습니다.

 

TENSOR RT INFERENCE

이제 INETWORK 에 weight를 일치시켜주어 필요한 Network 는 완성시켰습니다. 이제, Config 와 Network 를 통하여, builder_serialized_network() 에 넣어 serialized(한국말로 매칭되는 언어를 찾아주세요 ㅜㅜ) 시켜주면, TENSOR RT ENGINE이 완성됩니다. 

def build_engine(weights):
    '''
     생략 : 저번 포스팅에 적은 build engine 참고
    '''
	
    #runtime 즉, 실행되기 위해서 deserialize 하기 위한 객체
    runtime = trt.Runtime(TRT_LOGGER)
    
    # 네트워크 모델을 weight와 함께 생성
	populate_network(network, weights)
    
	plan = builder.build_serialized_network(network, config)
    return runtime.deserialize_cuda_engine(plan)
   
# 방법 2. 좀 더 직관적인 생산 방법
serialized_engine = builder.build_serialized_network(network, config)

공식문서에서의 build _ serialized _ network

 이렇게 builder.build_serialized_network 를 집어 넣주어, serialized network를 생성해줍니다. (복잡합니다.. 직관적이지 않지만 어쩔수 없습니다,,, C++ API 를 Wrapper 형태로만 감싼것이다 보니 약간은 어색할 수 있습니다.)

 

그럼 이제 serialized 된 network 를 저장하는 방법을 알아봅시다.

with open(“sample.engine”, “wb”) as f:
    f.write(serialized_engine)

다음과 같이 sample.engine 파일을 열어준다음, serialized_engine을 써주면 됩니다.

 

 

 

BUILD TENSOR RT RUNTIME

이렇게 완성된 네트워크를 runtime API 를 통해 deserialized 해주게 되면, 최종적으로 engine이 완성되어 모델을 실행할 수 있게됩니다. 다음과 같이 runtime을 통해 (여기서 runtime 또한 마찬가지로 TRT.LOGGER 를 넣어주어야함.) 이렇게 되면, 진정한 TENSOR RT engine이 만들어지게 됩니다.

runtime = trt.Runtime(logger)

engine = runtime.deserialize_cuda_engine(serialized_engine)

# 만약 위에서 처럼 serialized engine객체를 .engine으로 저장했을 경우,
with open("sample.engine","rb") as f:
	serialized_engine=f.read()

 

이제 engine이 만들어 졌으니, Inference 할 일만 남았습니다. 하지만 아직 한단계 남았습니다..(아직 한발 남았다..)

engine을 실행시키기 위한 engine.create_execution_context() 를 생성해주어야 합니다. (C++ 에서는 tensorrt.IExecutionContext 객체)

그다음 , context() 를 생성해주었으면, engine 에 input 과 output을 넣어주어야하는데, 이때, input 객체 와 output 객체 또한 gpu 상의 메모리 pointer를 가리켜 줘야 한다. 따라서, buffers 의 포인터를 생성해주어 GPU 상의 메모리를 할당해주고, 이렇게 할당된 buffers 를 context.execute_async_v2() 함수를 통해 inference 수행이 가능하다.

 

복잡하니 차근차근 다시 설명하자면, 생성된 engine을 바탕으로 context 생성

# TRT 엔진 생성
engine = runtime.deserialize_cuda_engine(serialized_engine)

# engine 을 실행시키기 위한, context 생성
context = engine.create_execution_context()

 

그 다음, 같은 GPU 상의  input, output 버퍼 생성.

# 버퍼 인덱스 생성
input_idx = engine[input_name]
output_idx = engine[output_name]

# GPU 포인터 생성 -> 파이썬은 포인터가 없지만 비슷하게 생성
buffers = [None] * 2 # Assuming 1 input and 1 output
buffers[input_idx] = input_ptr
buffers[output_idx] = output_ptr

 

마지막으로 CUDA STREAM 을 통한 추론 생성.

context.execute_async_v2(buffers, stream_ptr)

# 마지막으로 , 동기화를 위한 synchronize() 함수 실행
stream.synchronize()

 

 

이렇게 코드를 최종실행 할 수 있게 된다. 굉장히 많은 포인터들이 왔다갔다 하는데, 주의해서 코드를 짜길 바라겠습니다.. 다음 글을 참조 했으며, 기회가 되면 C++ API 또한 정리 할 수 있도록 하겠습니다 ㅎㅎ

https://github.com/NVIDIA/TensorRT/blob/master/samples/python/network_api_pytorch_mnist/sample.py  

GitHub - NVIDIA/TensorRT: TensorRT is a C++ library for high performance inference on NVIDIA GPUs and deep learning accelerators