# # Copyright 1993-2020 NVIDIA Corporation. All rights reserved. # # NOTICE TO LICENSEE: # # This source code and/or documentation ("Licensed Deliverables") are # subject to NVIDIA intellectual property rights under U.S. and # international Copyright laws. # # These Licensed Deliverables contained herein is PROPRIETARY and # CONFIDENTIAL to NVIDIA and is being provided under the terms and # conditions of a form of NVIDIA software license agreement by and # between NVIDIA and Licensee ("License Agreement") or electronically # accepted by Licensee. Notwithstanding any terms or conditions to # the contrary in the License Agreement, reproduction or disclosure # of the Licensed Deliverables to any third party without the express # written consent of NVIDIA is prohibited. # # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE # LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE # SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS # PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND. # NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED # DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY, # NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE. # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE # LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY # SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY # DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, # WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS # ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE # OF THESE LICENSED DELIVERABLES. # # U.S. Government End Users. These Licensed Deliverables are a # "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT # 1995), consisting of "commercial computer software" and "commercial # computer software documentation" as such terms are used in 48 # C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government # only as a commercial end item. Consistent with 48 C.F.R.12.212 and # 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all # U.S. Government End Users acquire the Licensed Deliverables with # only those rights set forth herein. # # Any use of the Licensed Deliverables in individual and commercial # software must include, in the user documentation and internal # comments to the code, the above Disclaimer and U.S. Government End # Users Notice. # import argparse import pycuda.driver as cuda import pycuda.autoinit import numpy as np import torch import tensorrt as trt import time import onnx import onnxruntime import os,sys,cv2 #from model.u2net import U2NET #cuda.init() model_names = ['u2net.onnx','u2net_dynamic_batch.onnx','u2net_dynamic_hw.onnx','u2net_dynamic_batch-hw.onnx' ] dynamic_batch = {'input':{0:'batch'}, 'output0':{0:'batch'}, 'output1':{0:'batch'}, 'output2':{0:'batch'}, 'output3':{0:'batch'}, 'output4':{0:'batch'}, 'output5':{0:'batch'}, 'output6':{0:'batch'}} dynamic_hw ={'input':{2:'H',3:'W'}, 'output0':{2:'H',3:'W'}, 'output1':{2:'H',3:'W'}, 'output2':{2:'H',3:'W'}, 'output3':{2:'H',3:'W'}, 'output4':{2:'H',3:'W'}, 'output5':{2:'H',3:'W'}, 'output6':{2:'H',3:'W'}} dynamic_batch_hw ={'input':{0:'batch',2:'H',3:'W'}, 'output0':{0:'batch',2:'H',3:'W'}, 'output1':{0:'batch',2:'H',3:'W'}, 'output2':{0:'batch',2:'H',3:'W'}, 'output3':{0:'batch',2:'H',3:'W'}, 'output4':{0:'batch',2:'H',3:'W'}, 'output5':{0:'batch',2:'H',3:'W'}, 'output6':{0:'batch',2:'H',3:'W'}} dynamic_=[None,dynamic_batch,dynamic_hw,dynamic_batch_hw] TRT_LOGGER = trt.Logger() def pth2onnx(pth_model,onnx_name,input_shape=(1,3,512,512),input_names=['input'],output_names=['output'],dynamix_axis=None): #pth_model:输入加载权重后的pth模型 #onnx_name:输出的onnx模型路径 #input_shape:模型输入的尺寸（建议尺寸） #input_names:模型输入的名字，list格式，可以有多个输入 #output_names:模型输入的名字，list格式，可以有多个输出 #dynamix_axis:字典格式，None-表示静态输入。每一个模型的输入输出都可以定义动态的维度 # 如dynamic_batch_hw ={'input':{0:'batch',2:'H',3:'W'}, 'output':{0:'batch',2:'H',3:'W'}}, # 表示input的B,H,W和output的B,H,W是动态尺寸 print('[I] beg to converting pth to onnx ...... ',dynamix_axis) input_tensor = torch.ones(input_shape) if next(pth_model.parameters()).is_cuda: input_tensor = input_tensor.to('cuda:0') with torch.no_grad(): torch.onnx.export(pth_model, input_tensor, onnx_name, opset_version=11, input_names=input_names, do_constant_folding=True, output_names=output_names, dynamic_axes=dynamix_axis) onnx_model = onnx.load(onnx_name) try: onnx.checker.check_model(onnx_model) except Exception as e: print('[Error] model incorrect:',e) else: print('[I] conver to onnx over in ', onnx_name) print('') def onnx_inference(onnx_input,model_name,outputName=['output0','output1','output2','output3','output4','output5','output6' ]): providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] print('8'*10, ' line125:',model_name) #outputName = ['pred_logits', 'pred_points'] onnx_session = onnxruntime.InferenceSession(model_name,providers=providers) try: onnx_output = onnx_session.run(outputName,onnx_input) except Exception as e: onnx_output=None print(e) return onnx_output def onnx2engine(onnx_file_path,engine_file_path,input_shape=[1,3,512,512],half=True,max_batch_size=1,input_profile_shapes=[None,None,None]): #onnx_file_path:输入的onnx路径 #engine_file_path:输出的trt模型路径 #input_shape:默认的模型输入尺寸, 如[1,3,512,512] ，如果是动态的可以为[1,3,-1,-1] #half:是否使用fp16,默认True #max_batch_size:最大的bachsize，默认是1 #input_profile_shapes:动态输入时输入的三个尺寸[最小尺寸,优化尺寸,最大尺寸]，此时input_shape一定有-1 # 如(1,3,512,512),(1,3,1024,1024),(1,3,2048,2048)， builder = trt.Builder(TRT_LOGGER) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) config = builder.create_builder_config() parser = trt.OnnxParser(network,TRT_LOGGER) runtime = trt.Runtime(TRT_LOGGER) # 最大内存占用,一般1G，trt特有的，一切与优化有关,显存溢出需要重新设置 config.max_workspace_size = 1<<30 #256MB if builder.platform_has_fast_fp16 and half: config.set_flag(trt.BuilderFlag.FP16) builder.max_batch_size = max_batch_size # 推理的时候要保证batch_size<=max_batch_size # parse model file if not os.path.exists(onnx_file_path): print(f'onnx file {onnx_file_path} not found,please run torch_2_onnx.py first to generate it') exit(0) print(f'Loading ONNX file from path {onnx_file_path}...') with open(onnx_file_path,'rb') as model: print('Beginning ONNX file parsing') 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 # Static input setting network.get_input(0).shape=input_shape # Dynamic input setting 动态输入在builder的profile设置 # 为每个动态输入绑定一个profile if -1 in input_shape: profile = builder.create_optimization_profile() profile.set_shape(network.get_input(0).name,input_profile_shapes[0],input_profile_shapes[1],input_profile_shapes[2] )#最小的尺寸,常用的尺寸,最大的尺寸,推理时候输入需要在这个范围内 config.add_optimization_profile(profile) print('Completed parsing the ONNX file') print(f'Building an engine from file {onnx_file_path}; this may take a while...') t0 = time.time() engine = builder.build_engine(network,config) with open(engine_file_path,'wb') as f: # f.write(plan) f.write(engine.serialize()) t1 = time.time() print('Completed creating Engine:%s, %.1f'%(engine_file_path,t1-t0)) try: # Sometimes python2 does not understand FileNotFoundError FileNotFoundError except NameError: FileNotFoundError = IOError #EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) def GiB(val): return val * 1 << 30 def add_help(description): parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter) args, _ = parser.parse_known_args() def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[]): ''' Parses sample arguments. Args: description (str): Description of the sample. subfolder (str): The subfolder containing data relevant to this sample find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path. Returns: str: Path of data directory. ''' # Standard command-line arguments for all samples. kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data") parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("-d", "--datadir", help="Location of the TensorRT sample data directory, and any additional data directories.", action="append", default=[kDEFAULT_DATA_ROOT]) args, _ = parser.parse_known_args() def get_data_path(data_dir): # If the subfolder exists, append it to the path, otherwise use the provided path as-is. data_path = os.path.join(data_dir, subfolder) if not os.path.exists(data_path): print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.") data_path = data_dir # Make sure data directory exists. if not (os.path.exists(data_path)): print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(data_path)) return data_path data_paths = [get_data_path(data_dir) for data_dir in args.datadir] return data_paths, locate_files(data_paths, find_files) def locate_files(data_paths, filenames): """ Locates the specified files in the specified data directories. If a file exists in multiple data directories, the first directory is used. Args: data_paths (List[str]): The data directories. filename (List[str]): The names of the files to find. Returns: List[str]: The absolute paths of the files. Raises: FileNotFoundError if a file could not be located. """ found_files = [None] * len(filenames) for data_path in data_paths: # Find all requested files. for index, (found, filename) in enumerate(zip(found_files, filenames)): if not found: file_path = os.path.abspath(os.path.join(data_path, filename)) if os.path.exists(file_path): found_files[index] = file_path # Check that all files were found for f, filename in zip(found_files, filenames): if not f or not os.path.exists(f): raise FileNotFoundError("Could not find {:}. Searched in data paths: {:}".format(filename, data_paths)) return found_files # Simple helper data class that's a little nicer to use than a 2-tuple. class HostDeviceMem(object): def __init__(self, host_mem, device_mem): self.host = host_mem self.device = device_mem def __str__(self): return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) def __repr__(self): return self.__str__() # Allocates all buffers required for an engine, i.e. host/device inputs/outputs. def allocate_buffers(engine,input_shape,streamFlag=True): inputs = [] outputs = [] bindings = [] if streamFlag: stream = cuda.Stream() else: stream=None for ib,binding in enumerate(engine): dims = engine.get_binding_shape(binding) #print(engine.get_binding_name(ib),dims,engine.max_batch_size) if -1 in dims: if isinstance(input_shape,list): dims = input_shape[ib] else: dims = input_shape # size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size #size = trt.volume(dims) * engine.max_batch_size size = trt.volume(dims) dtype = trt.nptype(engine.get_binding_dtype(binding)) # Allocate host and device buffers host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) # Append the device buffer to device bindings. bindings.append(int(device_mem)) # Append to the appropriate list. if engine.binding_is_input(binding): inputs.append(HostDeviceMem(host_mem, device_mem)) else: outputs.append(HostDeviceMem(host_mem, device_mem)) return inputs, outputs, bindings, stream # This function is generalized for multiple inputs/outputs. # inputs and outputs are expected to be lists of HostDeviceMem objects. def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): # Transfer input data to the GPU. [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] # Run inference. context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) # Transfer predictions back from the GPU. [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] # Synchronize the stream stream.synchronize() # Return only the host outputs. return [out.host for out in outputs] # This function is generalized for multiple inputs/outputs for full dimension networks. # inputs and outputs are expected to be lists of HostDeviceMem objects. def do_inference_v2(context, bindings, inputs, outputs, stream): # Transfer input data to the GPU. [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] # Run inference. #stream.synchronize() #context.execute_v2(bindings) # 执行推 context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) # Transfer predictions back from the GPU. [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] # Synchronize the stream stream.synchronize() # Return only the host outputs. return [out.host for out in outputs] def trt_inference( img,img_h,img_w,context,inputs,outputs,bindings,stream,input_name = 'input'): #输入： #img--np格式,NCHW #img_h,img_w--输入模型时图像的H，W。动态输入是需要知道。 #context--外面开辟的trt上下文 #inputs,outputs,bindings,stream--第一次处理图像时，开辟的内存及其地址绑定到trt的输出 #input_name--模型输入tensor的名字 #输出 #trt_outputs--为list格式，里面的元素是numpy格式 origin_inputshape = context.get_tensor_shape( input_name) #if origin_inputshape[-1]==-1: context.set_optimization_profile_async(0,stream.handle) origin_inputshape[-2],origin_inputshape[-1]=(img_h,img_w) context.set_input_shape(input_name, (origin_inputshape)) inputs[0].host = np.ascontiguousarray(img) trt_outputs = do_inference_v2(context,bindings=bindings,inputs=inputs,outputs=outputs,stream=stream) return trt_outputs # def do_inference_v3(context, bindings, inputs, outputs, stream,h_,w_): # ''' # Copy from https://github.com/zhaogangthu/keras-yolo3-ocr-tensorrt/blob/master/tensorRT_yolo3/common.py # # ''' # # Transfer input data to the GPU. # # context.set_binding_shape(0, (1, 3, h_, w_)) # # [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] # # Run inference. # context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) # # Transfer predictions back from the GPU. # [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] # # Synchronize the stream # stream.synchronize() # # Return only the host outputs. # return [out.host for out in outputs] if __name__=='__main__': model_path='weights/u2net_portrait.pth' onnx_name = model_path.replace('.pth','.onnx') trt_name = model_path.replace('.pth','.engine') pth_model = U2NET(3,1) pth_model.load_state_dict(torch.load(model_path)) input_names=['input'] output_names=['output%d'%(i) for i in range(7)] dynamix_axis = dynamic_hw input_shape =(1,3,512,512) #测试pth转为onnx模型 #pth2onnx(pth_model,onnx_name,input_shape=input_shape ,input_names=input_names ,output_names=output_names ,dynamix_axis=dynamix_axis ) #测试onnx模型转为trt模型 input_profile_shapes = [(1,3,512,512),(1,3,1024,1024),(1,3,2048,2048)] input_shape = [1,3,-1,-1] half=True max_batch_size = 1 onnx2engine(onnx_name,trt_name,input_shape=input_shape,half=half,max_batch_size=max_batch_size,input_profile_shapes=input_profile_shapes) ''' with torch.no_grad(): for i,model_name in enumerate(model_names): print(f'process model:{model_name}...') torch.onnx.export(model, input_tensor, model_name, opset_version=11, input_names=['input'], output_names=['output0','output1','output2','output3','output4','output5','output6'], dynamic_axes=dynamic_[i]) print(f'onnx model:{model_name} saved successfully...') #print('sleep 10s...') time.sleep(10) print(f'begin check onnx model:{model_name}...') onnx_model = onnx.load(model_name) try: onnx.checker.check_model(onnx_model) except Exception as e: print('model incorrect') print(e) else: print('model correct') print('*'*50) print('Begin to test...') case_1 = np.random.rand(1,3,512,512).astype(np.float32) case_2 = np.random.rand(2,3,512,512).astype(np.float32) case_3 = np.random.rand(1,3,224,224).astype(np.float32) cases = [case_1,case_2,case_3] providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] for model_name in model_names: print('-'*50,model_name) onnx_session = onnxruntime.InferenceSession(model_name,providers=providers) for i,case in enumerate(cases): onnx_input = {'input':case} try: onnx_output = onnx_session.run(['output0','output1','output2','output3','output4','output5','output6'],onnx_input)[0] except Exception as e: print(f'Input:{i} on model:{model_name} failed') print(e) else: print(f'Input:{i} on model:{model_name} succeed') '''