(batch_sizes, workspace_size, sequence_lengths, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose)
| 399 | return emb_layer |
| 400 | |
| 401 | def build_engine(batch_sizes, workspace_size, sequence_lengths, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose): |
| 402 | explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) |
| 403 | |
| 404 | with trt.Builder(TRT_LOGGER) as builder, builder.create_network(explicit_batch_flag) as network, builder.create_builder_config() as builder_config: |
| 405 | builder_config.max_workspace_size = workspace_size * (1024 * 1024) |
| 406 | builder_config.avg_timing_iterations = 8 |
| 407 | if config.use_fp16: |
| 408 | builder_config.set_flag(trt.BuilderFlag.FP16) |
| 409 | if config.use_int8: |
| 410 | builder_config.set_flag(trt.BuilderFlag.INT8) |
| 411 | if not config.use_qat: |
| 412 | calibrator = BertCalibrator(squad_json, vocab_file, calibrationCacheFile, 1, sequence_lengths[-1], calib_num) |
| 413 | builder_config.set_quantization_flag(trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION) |
| 414 | builder_config.int8_calibrator = calibrator |
| 415 | if config.use_strict: |
| 416 | builder_config.set_flag(trt.BuilderFlag.STRICT_TYPES) |
| 417 | |
| 418 | if verbose: |
| 419 | builder_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED |
| 420 | |
| 421 | if config.use_sparsity: |
| 422 | TRT_LOGGER.log(TRT_LOGGER.INFO, "Setting sparsity flag on builder_config.") |
| 423 | builder_config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS) |
| 424 | |
| 425 | # speed up the engine build for trt major version >= 8 |
| 426 | # 1. disable cudnn tactic |
| 427 | # 2. load global timing cache |
| 428 | if trt_version[0] >= 8: |
| 429 | tactic_source = builder_config.get_tactic_sources() & ~(1 << int(trt.TacticSource.CUDNN)) |
| 430 | builder_config.set_tactic_sources(tactic_source) |
| 431 | if config.timing_cache != None: |
| 432 | if os.path.exists(config.timing_cache): |
| 433 | with open(config.timing_cache, "rb") as f: |
| 434 | cache = builder_config.create_timing_cache(f.read()) |
| 435 | builder_config.set_timing_cache(cache, ignore_mismatch = False) |
| 436 | else: |
| 437 | cache = builder_config.create_timing_cache(b"") |
| 438 | builder_config.set_timing_cache(cache, ignore_mismatch = False) |
| 439 | |
| 440 | # only use the largest sequence when in calibration mode |
| 441 | if config.is_calib_mode: |
| 442 | sequence_lengths = sequence_lengths[-1:] |
| 443 | |
| 444 | # Create the network |
| 445 | emb_layer = emb_layernorm(builder, network, config, weights_dict, builder_config, sequence_lengths, batch_sizes) |
| 446 | embeddings = emb_layer.get_output(0) |
| 447 | mask_idx = emb_layer.get_output(1) |
| 448 | |
| 449 | bert_out = bert_model(config, weights_dict, network, embeddings, mask_idx) |
| 450 | |
| 451 | squad_logits = squad_output("cls_", config, weights_dict, network, bert_out) |
| 452 | squad_logits_out = squad_logits.get_output(0) |
| 453 | |
| 454 | network.mark_output(squad_logits_out) |
| 455 | |
| 456 | build_start_time = time.time() |
| 457 | engine = builder.build_engine(network, builder_config) |
| 458 | build_time_elapsed = (time.time() - build_start_time) |
no test coverage detected