很多開發人員在轉換完 TensorRT 加速引擎之后,最后準備調用起來執行推理任務的時候,就遇到一些障礙。這個環節是需要開發人員自行撰寫相關代碼,去執行讀入數據(前處理)、執行推理、顯示結果(后處理)等工作,如下圖最右邊的部分。
這部分的麻煩之處,在于每個神經網絡的結構不相同,并沒有“通用”的代碼可以適用于大部分的網絡結構,需要針對指定神經網絡去撰寫對應的代碼,最重要是需要清除這個模型的輸入 (input bold) 與輸出 (outpold) 的名稱與張量結構。
本文以前面在 TAO 工具套件中使用的 ssd 神經網絡為范例,提供基礎的“前后處理”范例代碼給讀者參考,這是從 NVIDIA 中國區開發者社區所舉辦過多屆 “Sky 黑客松”比賽中,所提供的開源內容中提取的重點,主要如下:
1、數據前處理:
def _preprocess_trt(img, shape=(300, 300)):
"""TRT SSD推理前的數據前處理"""
img = cv2.resize(img, shape)
img = img.transpose((2, 0, 1)).astype(np.float32)
returnimg
這里 “shape=(300,300)” 為張量的尺度,根據模型訓練時的長寬兩個變量,至于 transpose 里的 (2,0,1) 是固定的,不需調整。
2、數據后處理:
def _postprocess_trt(img, output, conf_th, output_layout):
"""TRT SSD推理后的結果的數據處理步驟."""
img_h, img_w, _ = img.shape
boxes, confs, clss = [], [], []
for prefix in range(0, len(output), output_layout):
index = int(output[prefix+0])
conf = float(output[prefix+2])
if conf < conf_th:
continue
x1 = int(output[prefix+3] * img_w)
y1 = int(output[prefix+4] * img_h)
x2 = int(output[prefix+5] * img_w)
y2 = int(output[prefix+6] * img_h)
cls = int(output[prefix+1])
boxes.append((x1, y1, x2, y2))
confs.append(conf)
clss.append(cls)
returnboxes,confs,clss#返回標框坐標、置信度、類別
這里最重要的 x1, y1,x2, y2 坐標值,必須根據 SSD 神經網絡所定義的規范去進行修改,其他部分可以通用于大部分神經網絡。
3、定義 TrtSSD 類封裝運行 TRT SSD 所需的東西:
class TrtSSD(object):
# 加載自定義組建,如果TRT版本小于7.0需要額外生成flattenconcat自定義組件庫
def _load_plugins(self):
if trt.__version__[0] < '7':
ctypes.CDLL("ssd/libflattenconcat.so")
trt.init_libnvinfer_plugins(self.trt_logger, '')
#加載通過Transfer Learning Toolkit生成的推理引擎
def _load_engine(self):
TRTbin = 'ssd/TRT_%s.bin' % self.model #請根據實際狀況自行修改
with open(TRTbin, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
return runtime.deserialize_cuda_engine(f.read())
#通過加載的引擎,生成可執行的上下文
def _create_context(self):
for binding in self.engine:
size = trt.volume(self.engine.get_binding_shape(binding)) *
self.engine.max_batch_size
##注意:這里的host_mem需要使用pagelockedmemory,以免內存被釋放
host_mem = cuda.pagelocked_empty(size, np.float32)
cuda_mem = cuda.mem_alloc(host_mem.nbytes)
self.bindings.append(int(cuda_mem))
if self.engine.binding_is_input(binding):
self.host_inputs.append(host_mem)
self.cuda_inputs.append(cuda_mem)
else:
self.host_outputs.append(host_mem)
self.cuda_outputs.append(cuda_mem)
return self.engine.create_execution_context()
# 初始化引擎
def __init__(self, model, input_shape, output_layout=7):
self.model = model
self.input_shape = input_shape
self.output_layout = output_layout
self.trt_logger = trt.Logger(trt.Logger.INFO)
self._load_plugins()
self.engine = self._load_engine()
self.host_inputs = []
self.cuda_inputs = []
self.host_outputs = []
self.cuda_outputs = []
self.bindings = []
self.stream = cuda.Stream()
self.context = self._create_context()
# 釋放引擎,釋放GPU顯存,釋放CUDA流
def __del__(self):
del self.stream
del self.cuda_outputs
del self.cuda_inputs
# 利用生成的可執行上下文執行推理
def detect(self, img, conf_th=0.3):
img_resized = _preprocess_trt(img, self.input_shape)
np.copyto(self.host_inputs[0], img_resized.ravel())
# 將處理好的圖片從CPU內存中復制到GPU顯存
cuda.memcpy_htod_async(
self.cuda_inputs[0], self.host_inputs[0], self.stream)
# 開始執行推理任務
self.context.execute_async(
batch_size=1,
bindings=self.bindings,
stream_handle=self.stream.handle)
# 將推理結果輸出從GPU顯存復制到CPU內存
cuda.memcpy_dtoh_async(
self.host_outputs[1], self.cuda_outputs[1], self.stream)
cuda.memcpy_dtoh_async(
self.host_outputs[0], self.cuda_outputs[0], self.stream)
self.stream.synchronize()
output = self.host_outputs[0]
return_postprocess_trt(img,output,conf_th,self.output_layout)
上面三個部分對不同神經網絡都是不同的內容,如果要參考 YOLO 神經網絡的對應內容,推薦參考https://github.com/jkjung-avt/tensorrt_demos開源項目,里面有完整的 YOLOv3 與 YOLOv4 的詳細內容。
本文的開源代碼可以在此鏈接下載完整的內容與配套的工具。
https://pan.baidu.com/s/1fGLBnzqtnRNpfD3PbileOA密碼: 99et
審核編輯 :李倩
-
神經網絡
+關注
關注
42文章
4774瀏覽量
100912 -
NVIDIA
+關注
關注
14文章
5025瀏覽量
103270
原文標題:NVIDIA Jetson Nano 2GB 系列文章(65):執行部署的 TensorRT 加速引擎
文章出處:【微信號:NVIDIA-Enterprise,微信公眾號:NVIDIA英偉達企業解決方案】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論