(builder, network, config, weights_dict, builder_config, max_sequence_length, batch_sizes)
| 353 | return OUT |
| 354 | |
| 355 | def emb_layernorm(builder, network, config, weights_dict, builder_config, max_sequence_length, batch_sizes): |
| 356 | input_ids = network.add_input(name="input_ids", dtype=trt.int32, shape=(-1,)) |
| 357 | segment_ids = network.add_input(name="segment_ids", dtype=trt.int32, shape=(-1,)) |
| 358 | cu_seqlens = network.add_input(name="cu_seqlens", dtype=trt.int32, shape=(-1,)) |
| 359 | max_seqlen = network.add_input(name="max_seqlen", dtype=trt.int32, shape=(-1,)) |
| 360 | |
| 361 | for batch_size in batch_sizes: |
| 362 | # Specify profiles |
| 363 | profile = builder.create_optimization_profile() |
| 364 | min_shape = (1,) |
| 365 | shape = (max_sequence_length*batch_size,) |
| 366 | profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape) |
| 367 | profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape) |
| 368 | profile.set_shape("cu_seqlens", min=min_shape, opt=(batch_size+1,), max=(batch_size+1,)) |
| 369 | profile.set_shape("max_seqlen", min=min_shape, opt=(max_sequence_length,), max=(max_sequence_length,)) |
| 370 | builder_config.add_optimization_profile(profile) |
| 371 | |
| 372 | wbeta = trt.PluginField("bert_embeddings_layernorm_beta", weights_dict["bert_embeddings_layernorm_beta"].numpy(), trt.PluginFieldType.FLOAT32) |
| 373 | wgamma = trt.PluginField("bert_embeddings_layernorm_gamma", weights_dict["bert_embeddings_layernorm_gamma"].numpy(), trt.PluginFieldType.FLOAT32) |
| 374 | wwordemb = trt.PluginField("bert_embeddings_word_embeddings", weights_dict["bert_embeddings_word_embeddings"].numpy(), trt.PluginFieldType.FLOAT32) |
| 375 | wtokemb = trt.PluginField("bert_embeddings_token_type_embeddings", weights_dict["bert_embeddings_token_type_embeddings"].numpy(), trt.PluginFieldType.FLOAT32) |
| 376 | wposemb = trt.PluginField("bert_embeddings_position_embeddings", weights_dict["bert_embeddings_position_embeddings"].numpy(), trt.PluginFieldType.FLOAT32) |
| 377 | output_fp16 = trt.PluginField("output_fp16", np.array([1 if config.use_fp16 or config.use_int8 else 0]).astype(np.int32), trt.PluginFieldType.INT32) |
| 378 | |
| 379 | pfc = trt.PluginFieldCollection([wbeta, wgamma, wwordemb, wtokemb, wposemb, output_fp16]) |
| 380 | fn = (emln_plg_creator3 if config.use_megatron else emln_plg_creator2).create_plugin("embeddings", pfc) |
| 381 | |
| 382 | inputs = [input_ids, segment_ids, cu_seqlens, max_seqlen] |
| 383 | emb_layer = network.add_plugin_v2(inputs, fn) |
| 384 | |
| 385 | if config.use_int8 and config.use_qat: |
| 386 | dr_input = weights_dict['l0_attention_self_query_input_amax'] |
| 387 | set_output_range(emb_layer, dr_input, out_idx=0) |
| 388 | |
| 389 | if config.use_megatron: |
| 390 | dr_skln1_res_in = weights_dict['l0_attention_output_add_residual_input_quantizer_amax'] |
| 391 | set_output_range(emb_layer, dr_skln1_res_in, out_idx=1) |
| 392 | |
| 393 | set_output_name(emb_layer, "embeddings_", "output") |
| 394 | return emb_layer, cu_seqlens, max_seqlen |
| 395 | |
| 396 | def build_engine(batch_sizes, workspace_size, sequence_length, config, weights_dict, squad_json, vocab_file, calibrationCacheFile, calib_num, verbose): |
| 397 | explicit_batch_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) |
no test coverage detected