📅  最后修改于: 2022-03-11 14:45:55.857000             🧑  作者: Mango
print('[info] Building an engine...')
if TRT_VERSION_MAJOR >= 7:
# Create an optimization profile (see Section 7.2 of https://docs.nvidia.com/deeplearning/sdk/pdf/TensorRT-Developer-Guide.pdf).
profile = builder.create_optimization_profile()
# FIXME: Hardcoded for ImageNet. The minimum/optimum/maximum dimensions of a dynamic input tensor are the same.
# profile.set_shape(input_tensor_name, (1, 3, 224, 224), (max_batch_size, 3, 224, 224), (max_batch_size, 3, 224, 224))
config = builder.create_builder_config()
config.add_optimization_profile(profile)
trt_model_engine = builder.build_engine(network, config)
trt_model_context = trt_model_engine.create_execution_context()
else:
# if trt version is below 7 then use the build_cuda_engine
trt_model_engine = builder.build_cuda_engine(network)
trt_model_context = trt_model_engine.create_execution_context()