def build_engine_from_onnx(onnx_path, engine_name, batch_size, TRT_LOGGER):
model = ModelProto()
with open(onnx_path, "rb") as f:
model.ParseFromString(f.read())
d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value
d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value
d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value
shape = [batch_size, d0, d1, d2]
engine = eng.build_engine(TRT_LOGGER, onnx_path, shape=shape)
eng.save_engine(engine, engine_name)
return engine
def main(args):
engine_name = args.plan_file
onnx_path = args.onnx_file
batch_size = 1
model = ModelProto()
with open(onnx_path, "rb") as f:
model.ParseFromString(f.read())
d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value
d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value
d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value
shape = [batch_size, d0, d1, d2]
engine = eng.build_engine(onnx_path, shape=shape)
eng.save_engine(engine, engine_name)
def create_engine(onnx_path, engine_output_path):
batch_size = 1
model = ModelProto()
with open(onnx_path, "rb") as f:
model.ParseFromString(f.read())
print('ONNX model laoded...')
print('Creating engine from this onnx file, ', onnx_path)
d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value
d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value
d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value
shape = [batch_size , d0, d1 ,d2]
engine = eng.build_engine(onnx_path, shape= shape)
eng.save_engine(engine, engine_output_path)
print('TRT engine created and saved at, ', engine_output_path)
def main(args):
engine_name = args.plan_file
onnx_path = args.onnx_file
batch_size = config.batch_size
model = ModelProto()
with open(onnx_path, "rb") as f:
print("parsing")
model.ParseFromString(f.read())
d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value
d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value
d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value
d3 = model.graph.input[0].type.tensor_type.shape.dim[4].dim_value
d4 = model.graph.input[0].type.tensor_type.shape.dim[5].dim_value
shape = [batch_size , d0, d1 , d2, d3, d4] # for CNN-LSTM model
#shape = [1, d0, d1 , d2, d3]
print("data shape required",shape)
engine = eng.build_engine(onnx_path, shape= shape)
eng.save_engine(engine, engine_name)
import engine as eng
import argparse
from onnx import ModelProto
import tensorrt as trt
engine_name = "yolo_pedestrian.plan"
onnx_path = "model.onnx"
batch_size = 1
model = ModelProto()
with open(onnx_path, "rb") as f:
model.ParseFromString(f.read())
d0 = model.graph.input[0].type.tensor_type.shape.dim[1].dim_value
d1 = model.graph.input[0].type.tensor_type.shape.dim[2].dim_value
d2 = model.graph.input[0].type.tensor_type.shape.dim[3].dim_value
shape = [batch_size, d0, d1, d2]
engine = eng.build_engine(onnx_path, shape=shape)
eng.save_engine(engine, engine_name)
import engine as eng
import argparse
from onnx import ModelProto
import tensorrt as trt
args = argparse.ArgumentParser()
args.add_argument('-engine-name')
args.add_argument('-onnx-path')
args = args.parse_args()
batch_size = 1
model = ModelProto()
with open(args.onnx_path, "rb") as f:
model.ParseFromString(f.read()) # load the model to get the input shape
dims = model.graph.input[0].type.tensor_type.shape.dim[1:4]
shape = [batch_size] + [d.dim_value for d in dims]
print(shape)
# build and save the engine
engine = eng.build_engine(args.onnx_path, shape)
eng.save_engine(engine, args.engine_name)