| 51 | |
| 52 | # Builds TensorRT Engine |
| 53 | def build_engine(model_path): |
| 54 | |
| 55 | builder = trt.Builder(TRT_LOGGER) |
| 56 | network = builder.create_network(common.EXPLICIT_BATCH) |
| 57 | config = builder.create_builder_config() |
| 58 | parser = trt.OnnxParser(network, TRT_LOGGER) |
| 59 | runtime = trt.Runtime(TRT_LOGGER) |
| 60 | |
| 61 | # Parse model file |
| 62 | print("Loading ONNX file from path {}...".format(model_path)) |
| 63 | with open(model_path, "rb") as model: |
| 64 | print("Beginning ONNX file parsing") |
| 65 | if not parser.parse(model.read()): |
| 66 | print("ERROR: Failed to parse the ONNX file.") |
| 67 | for error in range(parser.num_errors): |
| 68 | print(parser.get_error(error)) |
| 69 | return None |
| 70 | print("Completed parsing of ONNX file") |
| 71 | |
| 72 | config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1)) |
| 73 | |
| 74 | # The input text length is variable, so we need to specify an optimization profile. |
| 75 | profile = builder.create_optimization_profile() |
| 76 | for i in range(network.num_inputs): |
| 77 | input = network.get_input(i) |
| 78 | assert input.shape[0] == -1 |
| 79 | min_shape = [1] + list(input.shape[1:]) |
| 80 | opt_shape = [8] + list(input.shape[1:]) |
| 81 | max_shape = [MAX_TEXT_LENGTH] + list(input.shape[1:]) |
| 82 | profile.set_shape(input.name, min_shape, opt_shape, max_shape) |
| 83 | config.add_optimization_profile(profile) |
| 84 | |
| 85 | print("Building TensorRT engine. This may take a few minutes.") |
| 86 | plan = builder.build_serialized_network(network, config) |
| 87 | engine = runtime.deserialize_cuda_engine(plan) |
| 88 | with open(ENGINE_FILE_PATH, "wb") as f: |
| 89 | f.write(plan) |
| 90 | return engine |
| 91 | |
| 92 | def load_test_case(inputs, context_text, query_text, trt_context): |
| 93 | # Part 1: Specify Input shapes |