import keras2onnx
import onnx
from keras.models import load_model
model = load_model('testonnx.hdf5')
onnx_model = keras2onnx.convert_keras(model, model.name)
temp_model_file = 'result.onnx'
onnx.save_model(onnx_model, temp_model_file)

from PIL import Image
import numpy as np
import pycuda.driver as cuda
import time
import tensorrt as trt
import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import commonclass ModelData(object):MODEL_PATH = "result.onnx"INPUT_SHAPE = (1, 512, 512)# We can convert TensorRT data types to numpy types with trt.nptype()DTYPE = trt.float32# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)# Allocate host and device buffers, and create a stream.
def allocate_buffers(engine):# Determine dimensions and create page-locked memory buffers (i.e. won't be swapped to disk) to hold host inputs/outputs.h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(ModelData.DTYPE))h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(ModelData.DTYPE))# Allocate device memory for inputs and outputs.d_input = cuda.mem_alloc(h_input.nbytes)d_output = cuda.mem_alloc(h_output.nbytes)# Create a stream in which to copy inputs/outputs and run inference.stream = cuda.Stream()return h_input, d_input, h_output, d_output, streamdef do_inference(context, h_input, d_input, h_output, d_output, stream):# Transfer input data to the GPU.cuda.memcpy_htod_async(d_input, h_input, stream)# Run inference.context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)# Transfer predictions back from the GPU.cuda.memcpy_dtoh_async(h_output, d_output, stream)# Synchronize the streamstream.synchronize()# The Onnx path is used for Onnx models.
def build_engine_onnx(model_file):with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:builder.max_workspace_size = common.GiB(1)# Load the Onnx model and parse it in order to populate the TensorRT network.with open(model_file, 'rb') as model:parser.parse(model.read())last_layer = network.get_layer(network.num_layers - 1)network.mark_output(last_layer.get_output(0))returnresult=builder.build_cuda_engine(network)return returnresultdef load_normalized_test_case(test_image, pagelocked_buffer):# Converts the input image to a CHW Numpy arraydef normalize_image(image):# Resize, antialias and transpose the image to CHW.c, h, w = ModelData.INPUT_SHAPEimage_arr = np.asarray(image.resize((w, h), Image.ANTIALIAS))image_arr = np.reshape(image_arr, image_arr.shape + (1,))image_arr=image_arr.transpose([2, 0, 1])image_arr=image_arr.astype(trt.nptype(ModelData.DTYPE))image_arr=image_arr.ravel()# This particular ResNet50 model requires some preprocessing, specifically, mean normalization.return (image_arr / 255.0 - 0.45) / 0.225# Normalize the image and copy to pagelocked memory.np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))return test_imagedef main():  onnx_model_file='result.onnx'# Build a TensorRT engine.with build_engine_onnx(onnx_model_file) as engine:# Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same.# Allocate buffers and create a CUDA stream.h_input, d_input, h_output, d_output, stream = allocate_buffers(engine)with engine.create_execution_context() as context:# Load a normalized test case into the host input page-locked buffer.starttime=time.time()for i in range(100):test_image ='test.jpg'test_case = load_normalized_test_case(test_image, h_input)# Run the engine. The output will be a 1D tensor of length 1000, where each value represents the# probability that the image corresponds to that labeldo_inference(context, h_input, d_input, h_output, d_output, stream)#print('ok')endtime=time.time()pertime=(endtime-starttime)/100print('perimg cost'+str(pertime))

#修改后
def build_engine_onnx(model_file):with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:builder.max_workspace_size = common.GiB(1)# Load the Onnx model and parse it in order to populate the TensorRT network.with open(model_file, 'rb') as model:parser.parse(model.read())last_layer = network.get_layer(network.num_layers - 1)network.mark_output(last_layer.get_output(0))returnresult=builder.build_cuda_engine(network)return returnresult#原函数
def build_engine_onnx(model_file):with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:builder.max_workspace_size = common.GiB(1)# Load the Onnx model and parse it in order to populate the TensorRT network.with open(model_file, 'rb') as model:parser.parse(model.read())returnresult=builder.build_cuda_engine(network)return returnresult

last_layer = network.get_layer(network.num_layers - 1)
network.mark_output(last_layer.get_output(0))

[TensorRT] ERROR: Network must have at least one output