465 lines
20 KiB
Python
465 lines
20 KiB
Python
|
#
|
|||
|
|
# 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')
|
|||
|
|
'''
|