def build_engine(uff_model_path,
trt_logger,
trt_engine_datatype=trt.DataType.FLOAT,
batch_size=1,
silent=False):
#创建相关实例
with trt.Builder(trt_logger) as builder, builder.create_network(
) as network, builder.create_builder_config() as config, trt.UffParser(
) as parser:
#设置参数
config.max_workspace_size = 1 << 30
if trt_engine_datatype == trt.DataType.HALF:
config.set_flag(trt.BuilderFlag.FP16)
builder.max_batch_size = batch_size
#注册一个uff网络的输入节点和相关维度
#参考https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/parsers/Uff/pyUff.html?highlight=register_input#tensorrt.UffParser.register_input
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
#注册uff网络的输出节点
parser.register_output("MarkOutput_0")
#解析模型
parser.parse(uff_model_path, network)
if not silent:
print("Building TensorRT engine. This may take few minutes.")
#创建引擎
return builder.build_engine(network, config)
def build_engine(uff_model_path,
trt_logger,
trt_engine_datatype=trt.DataType.FLOAT,
calib_dataset=None,
batch_size=1,
silent=False):
with trt.Builder(trt_logger) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = 2 << 30
builder.max_batch_size = batch_size
if trt_engine_datatype == trt.DataType.HALF:
builder.fp16_mode = True
elif trt_engine_datatype == trt.DataType.INT8:
builder.fp16_mode = True
builder.int8_mode = True
builder.int8_calibrator = calibrator.SSDEntropyCalibrator(
data_dir=calib_dataset, cache_file='INT8CacheFile')
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output("MarkOutput_0")
parser.parse(uff_model_path, network)
if not silent:
print("Building TensorRT engine. This may take few minutes.")
return builder.build_cuda_engine(network)
def build_engine(uff_model_path,
calib,
trt_logger,
trt_engine_datatype=trt.DataType.FLOAT,
batch_size=1,
silent=False):
with trt.Builder(trt_logger) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = 1 << 30
if trt_engine_datatype == trt.DataType.HALF:
builder.fp16_mode = True
elif trt_engine_datatype == trt.DataType.INT8:
# builder.fp16_mode = True
builder.int8_mode = True
builder.int8_calibrator = calib
# builder.strict_type_constraints = True
builder.max_batch_size = batch_size
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output("MarkOutput_0")
parser.parse(uff_model_path, network)
if not silent:
print("Building TensorRT engine. This may take few minutes.")
return builder.build_cuda_engine(network)
def load(self, model_path):
uff_model = Path(model_path)
metadata_path = Path('%s/%s.metadata' %
(uff_model.parent.as_posix(), uff_model.stem))
with open(metadata_path.as_posix(), 'r') as metadata, \
trt.Builder(self.logger) as builder, \
builder.create_network() as network, \
trt.UffParser() as parser:
builder.max_workspace_size = 1 << 20
builder.max_batch_size = 1
builder.fp16_mode = True
metadata = json.loads(metadata.read())
# Configure inputs and outputs
print('Configuring I/O')
input_names = metadata['input_names']
output_names = metadata['output_names']
for name in input_names:
parser.register_input(
name,
(self.cfg.TARGET_D, self.cfg.TARGET_H, self.cfg.TARGET_W))
for name in output_names:
parser.register_output(name)
# Parse network
print('Parsing TensorRT Network')
parser.parse(uff_model.as_posix(), network)
print('Building CUDA Engine')
self.engine = builder.build_cuda_engine(network)
# Allocate buffers
print('Allocating Buffers')
self.inputs, self.outputs, self.bindings, self.stream \
= TensorRTLinear.allocate_buffers(self.engine)
print('Ready')
def build(self,
uff_model_path,
trt_logger,
trt_engine_datatype=trt.DataType.FLOAT,
calib_dataset=None,
batch_size=1):
print('Building TensorRT engine. This may take few minutes.')
print('uff_model_path:', uff_model_path)
print('trt_engine_datatype:', trt_engine_datatype)
with trt.Builder(trt_logger) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = 2 << 30 # 1GiB
self.max_batch_size = batch_size
builder.max_batch_size = batch_size
if trt_engine_datatype == trt.DataType.HALF:
builder.fp16_mode = True
elif trt_engine_datatype == trt.DataType.INT8:
builder.fp16_mode = True
builder.int8_mode = True
builder.int8_calibrator = Calibrator(
data_dir=calib_dataset, cache_file='INT8CacheFile')
parser.register_input(DetectionModel.input_name,
DetectionModel.input_shape)
parser.register_output(
DetectionModel.output_name) # "MarkOutput_0" ???
parser.parse(uff_model_path, network)
self.engine = builder.build_cuda_engine(network)
def parse_uff(self, uff_file, input_names, input_shapes, output_names):
"""Parses .uff file and prepares for serialization
:param uff_file: path to uff model
:param input_names: names of input as list
:param input_shapes: input shape (channels first)
:param output_names: names of output
"""
assert not isinstance(input_names, str) and not isinstance(
output_names, str), "supply I/O as lists"
parser = trt.UffParser()
for input_name, input_shape in zip(input_names, input_shapes):
# input shape must always be channels-first
parser.register_input(input_name, input_shapes)
for output_name in output_names:
parser.register_output(output_name)
parser.parse(uff_file, self.network,
CudaEngineManager.CONSTANTS["dtype"])
self.parser = parser
def main():
# Set the data path to the directory that contains the trained models and test images for inference.
#data_path, data_files = common.find_sample_data(description="Runs a ResNet50 network with a TensorRT inference engine.", subfolder="resnet50", find_files=["binoculars.jpeg", "reflex_camera.jpeg", "tabby_tiger_cat.jpg", ModelData.MODEL_PATH, "class_labels.txt"])
# Get test images, models and labels.
#test_images = data_files[0:3]
#uff_model_file, labels_file = data_files[3:]
#labels = open(labels_file, 'r').read().split('\n')
#uff_model_file = "/host_temp/pose-estimation-trt/data/CPN_TENSORRT_FP16.uff"
#model_file = "/host_temp/pose-estimation-trt/data/CPN_TENSORRT_FP16.uff"
#model_file = "/host_temp/pose-estimation-trt/data/CPN_TENSORFLOW_FROZEN.uff"
model_file = "/usr/src/tensorrt/samples/python/introductory_parser_samples/CPN_TENSORFLOW_FROZEN.uff"
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
parser.register_input("tower_0/Placeholder", (32, 256, 192, 3))
parser.register_output("tower_0/refine_out/BatchNorm/FusedBatchNorm")
# uff wrong!!!
parser.parse(model_file, network)
# error Segmentation fault (core dumped) -> parser location
#[TensorRT] ERROR: UFFParser: Validator error: TRTEngineOp_1: Unsupported operation _TRTEngineOp
'''
def main():
parser = argparse.ArgumentParser()
parser.add_argument('model', type=str, choices=list(MODEL_SPECS.keys()))
args = parser.parse_args()
# initialize
if trt.__version__[0] < '7':
ctypes.CDLL(LIB_FILE)
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
# compile the model into TensorRT engine
model = args.model
spec = MODEL_SPECS[model]
dynamic_graph = add_plugin(gs.DynamicGraph(spec['input_pb']), model, spec)
_ = uff.from_tensorflow(dynamic_graph.as_graph_def(),
output_nodes=['NMS'],
output_filename=spec['tmp_uff'],
text=True,
debug_mode=DEBUG_UFF)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = 1 << 28
builder.max_batch_size = 1
builder.fp16_mode = True
parser.register_input('Input', INPUT_DIMS)
parser.register_output('MarkOutput_0')
parser.parse(spec['tmp_uff'], network)
engine = builder.build_cuda_engine(network)
buf = engine.serialize()
with open(spec['output_bin'], 'wb') as f:
f.write(buf)
def __call__(self):
uff_model, input_names, input_shapes, output_names = self.uff_loader()
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network()
parser = trt.UffParser()
# Input names should come from the converter, as a preprocessing script may have been applied to the frozen model.
for name, shape in zip(input_names, input_shapes):
# Default order is NCHW, only set to NHWC if we're reasonably certain that it is.
input_order = self.uff_order
if not self.uff_order:
input_order = trt.UffInputOrder.NCHW
if FormatManager.determine_format(shape) == DataFormat.NHWC:
input_order = trt.UffInputOrder.NHWC
shape = shape[1:]
G_LOGGER.verbose(
"Registering UFF input: {:} with shape: {:} and input order: {:}"
.format(name, shape, input_order))
parser.register_input(name, shape, input_order)
if output_names and output_names != constants.MARK_ALL:
for name in output_names:
G_LOGGER.verbose("Registering UFF output: " + str(name))
parser.register_output(name)
G_LOGGER.info(
"Parsing UFF model with inputs: {:} and outputs: {:}".format(
input_names, output_names))
success = parser.parse_buffer(uff_model, network)
if not success:
G_LOGGER.critical("Could not parse UFF correctly")
return builder, network, parser, input_shapes[0][0]
def build_save_engine(model_file):
# For more information on TRT basics, refer to the introductory samples.
engine_file_path='Body.engine'
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
builder.max_workspace_size = 1 << 30
# default FP32
# FP16
#builder.fp16_mode = True
# Parse the Uff Network
parser.register_input("input_1", [3, 224, 224], order=trt.UffInputOrder.NCHW)
parser.register_output("fc3/Softmax")
parser.parse(model_file, network)
print('parser done')
# Build and return an engine.
with builder.build_cuda_engine(network) as engine:
print('engine done')
with open(engine_file_path, "wb") as f:
f.write(engine.serialize())
def main():
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
# compile the model into TensorRT engine
model = 'ssd_mobilenet_v2_coco'
spec = MODEL_SPECS[model]
if not os.path.exists(spec['tmp_uff']):
dynamic_graph = add_plugin(gs.DynamicGraph(spec['input_pb']), spec)
uff_model = uff.from_tensorflow(dynamic_graph.as_graph_def(),
output_nodes=['NMS'],
output_filename=spec['tmp_uff'],
text=True,
debug_mode=DEBUG_UFF)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = 1 << 28
builder.max_batch_size = 1
builder.fp16_mode = True
parser.register_input('Input', INPUT_DIMS)
parser.register_output('MarkOutput_0')
parser.parse(spec['tmp_uff'], network)
print("Building Tensorrt engine. This may take a few minutes.")
engine = builder.build_cuda_engine(network)
buf = engine.serialize()
with open(spec['output_bin'], 'wb') as f:
f.write(buf)
print("Save engine.")
def initialize(self):
# Create network.
self.network = self.builder.create_network()
# Do graph surgery on pb graph and convert to UFF.
uff_model = uff.from_tensorflow_frozen_model(
self.model_path,
preprocessor="code/ssd-mobilenet/tensorrt/SSDMobileNet.py")
# Parse UFF model and populate network.
parser = trt.UffParser()
parser.register_input("Input", [3, 300, 300], trt.UffInputOrder.NCHW)
parser.register_output("Postprocessor")
success = parser.parse_buffer(uff_model, self.network)
if not success:
raise RuntimeError("SSDMobileNet network creation failed!")
# Set input dtype and format
input_tensor = self.network.get_input(0)
if self.input_dtype == "int8":
input_tensor.dtype = trt.int8
input_tensor.dynamic_range = (-1.0, 1.0)
if self.input_format == "linear":
input_tensor.allowed_formats = 1 << int(trt.TensorFormat.LINEAR)
elif self.input_format == "chw4":
input_tensor.allowed_formats = 1 << int(trt.TensorFormat.CHW4)
self.postprocess(replace_relu6=(self.dla_core is not None))
self.initialized = True
def verify_generated_uff_with_tensorrt_uffparser(ufffilename):
""" Loads the UFF file with TensorRT (py). """
assert os.path.isfile(
ufffilename), "ufffilename=" + ufffilename + ' doesnt exist'
import tensorrt as trt
print tcol.HEADER, '[verify_generated_uff_with_tensorrt_uffparser] TensorRT version=', trt.__version__, tcol.ENDC
try:
uffinput = "input_1"
uffinput_dims = (3, 240, 320)
# uffinput_dims = (256, 80,60)
uffoutput = "conv_pw_5_relu/Relu6"
# uffoutput = "net_vlad_layer_1/l2_normalize_1"
# uffoutput = "net_vlad_layer_1/add_1"
# uffoutput = "net_vlad_layer_1/Reshape_1"
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
print 'ufffilename=', str(ufffilename)
print 'uffinput=', str(uffinput), '\t', 'uffinput_dims=', str(
uffinput_dims)
print 'uffoutput=', str(uffoutput)
parser.register_input(uffinput, uffinput_dims)
parser.register_output(uffoutput)
parser.parse(ufffilename, network)
pass
print tcol.OKGREEN, '[verify_generated_uff_with_tensorrt_uffparser] Verified.....!', tcol.ENDC
except:
print tcol.FAIL, '[verify_generated_uff_with_tensorrt_uffparser] UFF file=', ufffilename, ' with uffinput=', uffinput, ' uffoutput=', uffoutput, ' cannot be parsed.'
def build_engine_uff(self):
with trt.Builder(self.trt_logger) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = self.GiB(1)
parser.register_input(self.model_input_name, self.input_shape)
parser.register_output(self.model_output_name)
parser.parse(self.model_path, network)
return builder.build_cuda_engine(network)
def __init__(self, file_path, input_name, input_shape, output_name, uff_buffer=None):
super().__init__(trt.UffParser())
self.file_path = file_path
self.input_name = input_name
self.input_shape = input_shape
self.output_name = output_name
self.uff_buffer = uff_buffer
def export_trt(pb_file, output_dir, num_classes=90, neuralet_adaptive_model=1):
"""
Exports the Tensorflow pb models to TensorRT engines.
Args:
pb_file: The path of input pb file
output_dir: A directory to store the output files
num_classes: Detector's number of classes
"""
lib_flatten_concat_file = "exporters/libflattenconcat.so.6"
# initialize
if trt.__version__[0] < '7':
ctypes.CDLL(lib_flatten_concat_file)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
# compile the model into TensorRT engine
model = "ssd_mobilenet_v2_coco"
if not os.path.isfile(pb_file):
raise FileNotFoundError(
'model does not exist under: {}'.format(pb_file))
if not os.path.isdir(output_dir):
print("the provided output directory : {0} is not exist".format(
output_dir))
print("creating output directory : {0}".format(output_dir))
os.makedirs(output_dir, exist_ok=True)
dynamic_graph = plugin.add_plugin_and_preprocess(gs.DynamicGraph(pb_file),
model, num_classes,
neuralet_adaptive_model)
model_file_name = ".".join((pb_file.split("/")[-1]).split(".")[:-1])
uff_path = os.path.join(output_dir, model_file_name + ".uff")
_ = uff.from_tensorflow(dynamic_graph.as_graph_def(),
output_nodes=['NMS'],
output_filename=uff_path,
text=True,
debug_mode=False)
input_dims = (3, 300, 300)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, builder.create_builder_config(
) as builder_config, trt.UffParser() as parser:
builder_config.max_workspace_size = 1 << 28
builder.max_batch_size = 1
builder_config.set_flag(trt.BuilderFlag.FP16)
parser.register_input('Input', input_dims)
parser.register_output('MarkOutput_0')
parser.parse(uff_path, network)
engine = builder.build_engine(network, builder_config)
buf = engine.serialize()
engine_path = os.path.join(output_dir, model_file_name + ".bin")
with open(engine_path, 'wb') as f:
f.write(buf)
print(
"your model has been converted to trt engine successfully under : {}"
.format(engine_path))
def build_engine(model_data):
uff_model = uff.from_tensorflow_frozen_model(model_data.pb_file_path)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
builder.max_workspace_size = GiB(5)
builder.fp16_mode = model_data.fp16_mode
parser.register_input(model_data.input_name, model_data.input_shape)
parser.register_output(model_data.output_name)
parser.parse_buffer(uff_model, network)
return builder.build_cuda_engine(network)
def build_engine(uff_path): with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser: builder.max_workspace_size = GiB(1) parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE) parser.register_output(ModelData.OUTPUT_NAME) parser.parse(uff_path, network) return builder.build_cuda_engine(network)
def convert_tf_model_to_trt(tf_model_filename, trt_model_filename,
model_data_layout, input_layer_name, input_height,
input_width, output_layer_name, output_data_type,
max_workspace_size, max_batch_size):
"Convert an tf_model_filename into a trt_model_filename using the given parameters"
uff_model = uff.from_tensorflow_frozen_model(tf_model_filename)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
if model_data_layout == 'NHWC':
parser.register_input(input_layer_name,
[input_height, input_width, 3],
trt.UffInputOrder.NHWC)
else:
parser.register_input(input_layer_name,
[3, input_height, input_width],
trt.UffInputOrder.NCHW)
parser.register_output(output_layer_name)
if not parser.parse_buffer(uff_model, network):
raise RuntimeError(
"UFF model parsing (originally from {}) failed. Error: {}".
format(tf_model_filename,
parser.get_error(0).desc()))
if (output_data_type == 'fp32'):
print('Converting into fp32 (default), max_batch_size={}'.format(
max_batch_size))
else:
if not builder.platform_has_fast_fp16:
print(
'Warning: This platform is not optimized for fast fp16 mode'
)
builder.fp16_mode = True
print('Converting into fp16, max_batch_size={}'.format(
max_batch_size))
builder.max_workspace_size = max_workspace_size
builder.max_batch_size = max_batch_size
trt_model_object = builder.build_cuda_engine(network)
try:
serialized_trt_model = trt_model_object.serialize()
with open(trt_model_filename, "wb") as trt_model_file:
trt_model_file.write(serialized_trt_model)
except:
raise RuntimeError(
'Cannot serialize or write TensorRT engine to file {}.'.format(
trt_model_filename))
def build_engine(model_file):
# For more information on TRT basics, refer to the introductory samples.
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
builder.max_workspace_size = common.GiB(1)
# Parse the Uff Network
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output(ModelData.OUTPUT_NAME)
parser.parse(model_file, network)
# Build and return an engine.
return builder.build_cuda_engine(network)
def build_engine(model_file, TRT_LOGGER):
# For more information on TRT basics, refer to the introductory samples.
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = 1 << 16
# Parse the Uff Network
parser.register_input("input_1", (3, 150, 150))
parser.register_output('dense_2/Softmax')
parser.parse(model_file, network)
# Build and return an engine.
return builder.build_cuda_engine(network)
def build_engine(model_file, logger, data):
# For more information on TRT basics, refer to the introductory samples.
with trt.Builder(logger) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_workspace_size = GiB(2)
builder.max_batch_size = 10
# Parse the Uff Network
parser.register_input(data["INPUT_NAME"], data["INPUT_SHAPE"])
parser.register_output("Openpose/concat_stage7")
parser.parse(model_file, network)
# Build and return an engine.
return builder.build_cuda_engine(network)
def build_engine(model_file):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
builder.max_batch_size = trt_batch_size
builder.max_workspace_size = common.GiB(
1) # Workspace size是builder在构建engine时候最大可以使用的内存大小，其越高越好
# parser the UFF Nerwork
parser.register_input('input_1', (3, 224, 224))
parser.register_output("fc1000/Softmax")
parser.parse(model_file, network) # 载入模型，解析，填充tensorRT的network
print "build engine..."
return builder.build_cuda_engine(network) # build network
def build_engine(model_file):
print('build engine...')
with trt.Builder(LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
builder.max_workspace_size = MAX_WORKSPACE_SIZE
builder.max_batch_size = MAX_BATCH_SIZE
if DTYPE == trt.float16:
builder.fp16_mode = True
parser.register_input(INPUT_NAME, INPUT_SHAPE, trt.UffInputOrder.NCHW)
parser.register_output(OUTPUT_NAME)
parser.parse(model_file, network, DTYPE)
return builder.build_cuda_engine(network)
def build_engine(model_file):
# For more information on TRT basics, refer to the introductory samples.
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, builder.create_builder_config() as config, trt.UffParser(
) as parser, trt.Runtime(TRT_LOGGER) as runtime:
config.max_workspace_size = common.GiB(1)
# Parse the Uff Network
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output(ModelData.OUTPUT_NAME)
parser.parse(model_file, network)
# Build and return an engine.
plan = builder.build_serialized_network(network, config)
return runtime.deserialize_cuda_engine(plan)
def build_engine_uff(model_file):
# You can set the logger severity higher to suppress messages (or lower to display more messages).
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
# Workspace size is the maximum amount of memory available to the builder while building an engine.
# It should generally be set as high as possible.
builder.max_workspace_size = common.GiB(1)
# We need to manually register the input and output nodes for UFF.
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output(ModelData.OUTPUT_NAME)
# Load the UFF model and parse it in order to populate the TensorRT network.
parser.parse(model_file, network)
# Build and return an engine.
return builder.build_cuda_engine(network)
def build_engine(model_path):
#创建相应的实例
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, builder.create_builder_config() as config, trt.UffParser(
) as parser:
#配置相关参数
config.max_workspace_size = common.GiB(1)
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output(ModelData.OUTPUT_NAME)
parser.parse(model_path, network)
#构建引擎
return builder.build_engine(network, config)
def build_engine(self):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, trt.UffParser() as parser:
parser.register_input("input", (1, 1, 1024))
parser.register_output("generator/Tanh")
parser.parse(self.uff_file_name, network)
builder.max_workspace_size = 1 << 20
builder.fp16_mode = True
builder.strict_type_constraints = True
engine = builder.build_cuda_engine(network)
return engine
def build_int8_engine(model_file_path, calib, data, logger):
with trt.Builder(logger) as builder, \
builder.create_network() as network, \
trt.UffParser() as parser:
builder.max_batch_size = int(os.environ["BATCH_SIZE"])
builder.max_workspace_size = common.GiB(2)
builder.int8_mode = True
builder.int8_calibrator = calib
parser.register_input(data["INPUT_NAME"], data["INPUT_SHAPE"])
parser.register_output("Openpose/concat_stage7")
parser.parse(model_file_path, network)
# Build and return an engine.
engine = builder.build_cuda_engine(network)
return engine
def build_engine(model_path):
""" Create TRT Engine from uff with custom layer plugins """
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
) as network, builder.create_builder_config() as config, trt.UffParser(
) as parser:
config.max_workspace_size = common.GiB(1)
uff_path = model_to_uff(model_path)
print("pb to uff is done")
parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
parser.register_output(ModelData.OUTPUT_NAME)
parser.parse(uff_path, network)
return builder.build_engine(network, config)
