tensorrt動態(tài)輸入分辨率尺寸

本文只有 tensorrt python部分涉動態(tài)分辨率設(shè)置，沒有c++的。

目錄

pytorch轉(zhuǎn)onnx：

onnx轉(zhuǎn)tensorrt：

python tensorrt推理：

知乎博客也可以參考：

tensorrt動態(tài)輸入（Dynamic shapes） - 知乎

記錄此貼的原因有兩個：1.肯定也有很多人需要 。2.就我搜索的帖子沒一個講的明明白白的，官方文檔也不利索，需要連蒙帶猜。話不多少，直接上代碼。

以pytorch轉(zhuǎn)onnx轉(zhuǎn)tensorrt為例，動態(tài)shape是圖像的長寬。

pytorch轉(zhuǎn)onnx：

def export_onnx(model,image_shape,onnx_path, batch_size=1):

x,y=image_shape

img = torch.zeros((batch_size, 3, x, y))

dynamic_onnx=True

if dynamic_onnx:

dynamic_ax = {'input_1' : {2 : 'image_height',3:'image_wdith'},

'output_1' : {2 : 'image_height',3:'image_wdith'}}

torch.onnx.export(model, (img), onnx_path,

input_names=["input_1"], output_names=["output_1"], verbose=False, opset_version=11,dynamic_axes=dynamic_ax)

else:

torch.onnx.export(model, (img), onnx_path,

input_names=["input_1"], output_names=["output_1"], verbose=False, opset_version=11

)

onnx轉(zhuǎn)tensorrt：

按照nvidia官方文檔對dynamic shape的定義，所謂動態(tài)，無非是定義engine的時候不指定，用-1代替，在推理的時候再確定，因此建立engine 和推理部分的代碼都需要修改。

建立engine時，從onnx讀取的network，本身的輸入輸出就是dynamic shapes，只需要增加optimization_profile來確定一下輸入的尺寸范圍。

def build_engine(onnx_path, using_half,engine_file,dynamic_input=True):

trt.init_libnvinfer_plugins(None, '')

with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:

builder.max_batch_size = 1 # always 1 for explicit batch

config = builder.create_builder_config()

config.max_workspace_size = GiB(1)

if using_half:

config.set_flag(trt.BuilderFlag.FP16)

# Load the Onnx model and parse it in order to populate the TensorRT network.

with open(onnx_path, 'rb') as model:

if not parser.parse(model.read()):

print ('ERROR: Failed to parse the ONNX file.')

for error in range(parser.num_errors):

print (parser.get_error(error))

return None

##增加部分

if dynamic_input:

profile = builder.create_optimization_profile();

profile.set_shape("input_1", (1,3,512,512), (1,3,1024,1024), (1,3,1600,1600))

config.add_optimization_profile(profile)

#加上一個sigmoid層

previous_output = network.get_output(0)

network.unmark_output(previous_output)

sigmoid_layer=network.add_activation(previous_output,trt.ActivationType.SIGMOID)

network.mark_output(sigmoid_layer.get_output(0))

return builder.build_engine(network, config)

python tensorrt推理：

進行推理時，有個不小的暗坑，按照我之前的理解，既然動態(tài)輸入，我只需要在給輸入分配合適的緩存，然后不管什么尺寸直接推理就行了唄，事實證明還是年輕了。按照官方文檔的提示，在推理的時候一定要增加這么一行，context.active_optimization_profile = 0，來選擇對應(yīng)的optimization_profile，ok，我加了，但是還是報錯了，原因是我們既然在定義engine的時候沒有定義輸入尺寸，那么在推理的時候就需要根據(jù)實際的輸入定義好輸入尺寸。

def profile_trt(engine, imagepath,batch_size):

assert(engine is not None)

input_image,input_shape=preprocess_image(imagepath)

segment_inputs, segment_outputs, segment_bindings = allocate_buffers(engine, True,input_shape)

stream = cuda.Stream()

with engine.create_execution_context() as context:

context.active_optimization_profile = 0#增加部分

origin_inputshape=context.get_binding_shape(0)

#增加部分

if (origin_inputshape[-1]==-1):

origin_inputshape[-2],origin_inputshape[-1]=(input_shape)

context.set_binding_shape(0,(origin_inputshape))

input_img_array = np.array([input_image] * batch_size)

img = torch.from_numpy(input_img_array).float().numpy()

segment_inputs[0].host = img

[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in segment_inputs]#Copy from the Python buffer src to the device pointer dest (an int or a DeviceAllocation) asynchronously,

stream.synchronize()#Wait for all activity on this stream to cease, then return.

context.execute_async(bindings=segment_bindings, stream_handle=stream.handle)#Asynchronously execute inference on a batch.

stream.synchronize(）

[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in segment_outputs]#Copy from the device pointer src (an int or a DeviceAllocation) to the Python buffer dest asynchronously

stream.synchronize()

results = np.array(segment_outputs[0].host).reshape(batch_size, input_shape[0],input_shape[1])

return results.transpose(1,2,0)

pytorch

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔相應(yīng)法律責任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實的內(nèi)容，請聯(lián)系我們jiasou666@gmail.com 處理，核實后本網(wǎng)站將在24小時內(nèi)刪除侵權(quán)內(nèi)容。