(builder, network, config, weights_dict, builder_config, sequence_lengths, batch_sizes)
| 342 | return OUT |
| 343 | |
| 344 | def emb_layernorm(builder, network, config, weights_dict, builder_config, sequence_lengths, batch_sizes): |
| 345 | # int8 only support some of the sequence length, we dynamic on sequence length is not allowed. |
| 346 | input_ids = network.add_input(name="input_ids", dtype=trt.int32, shape=(-1 if len(batch_sizes) > 1 else batch_sizes[0], -1 if len(sequence_lengths) > 1 else sequence_lengths[0])) |
| 347 | segment_ids = network.add_input(name="segment_ids", dtype=trt.int32, shape=(-1 if len(batch_sizes) > 1 else batch_sizes[0], -1 if len(sequence_lengths) > 1 else sequence_lengths[0])) |
| 348 | input_mask = network.add_input(name="input_mask", dtype=trt.int32, shape=(-1 if len(batch_sizes) > 1 else batch_sizes[0], -1 if len(sequence_lengths) > 1 else sequence_lengths[0])) |
| 349 | |
| 350 | # Specify profiles for the batch sizes we're interested in. |
| 351 | # Make sure the profile also works for all sizes not covered by the previous profile. |
| 352 | |
| 353 | if len(sequence_lengths) > 1 or len(batch_sizes) > 1: |
| 354 | for batch_size in sorted(batch_sizes): |
| 355 | if len(sequence_lengths) == 1: |
| 356 | profile = builder.create_optimization_profile() |
| 357 | min_shape = (1, sequence_lengths[0]) |
| 358 | shape = (batch_size, sequence_lengths[0]) |
| 359 | profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape) |
| 360 | profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape) |
| 361 | profile.set_shape("input_mask", min=min_shape, opt=shape, max=shape) |
| 362 | builder_config.add_optimization_profile(profile) |
| 363 | else: |
| 364 | for sequence_length in sorted(sequence_lengths): |
| 365 | profile = builder.create_optimization_profile() |
| 366 | min_shape = (1, sequence_length) |
| 367 | shape = (batch_size, sequence_length) |
| 368 | profile.set_shape("input_ids", min=min_shape, opt=shape, max=shape) |
| 369 | profile.set_shape("segment_ids", min=min_shape, opt=shape, max=shape) |
| 370 | profile.set_shape("input_mask", min=min_shape, opt=shape, max=shape) |
| 371 | builder_config.add_optimization_profile(profile) |
| 372 | |
| 373 | wbeta = trt.PluginField("bert_embeddings_layernorm_beta", weights_dict["bert_embeddings_layernorm_beta"].numpy(), trt.PluginFieldType.FLOAT32) |
| 374 | wgamma = trt.PluginField("bert_embeddings_layernorm_gamma", weights_dict["bert_embeddings_layernorm_gamma"].numpy(), trt.PluginFieldType.FLOAT32) |
| 375 | wwordemb = trt.PluginField("bert_embeddings_word_embeddings", weights_dict["bert_embeddings_word_embeddings"].numpy(), trt.PluginFieldType.FLOAT32) |
| 376 | wtokemb = trt.PluginField("bert_embeddings_token_type_embeddings", weights_dict["bert_embeddings_token_type_embeddings"].numpy(), trt.PluginFieldType.FLOAT32) |
| 377 | wposemb = trt.PluginField("bert_embeddings_position_embeddings", weights_dict["bert_embeddings_position_embeddings"].numpy(), trt.PluginFieldType.FLOAT32) |
| 378 | |
| 379 | output_fp16 = trt.PluginField("output_fp16", np.array([1 if config.use_fp16 else 0]).astype(np.int32), trt.PluginFieldType.INT32) |
| 380 | mha_type = trt.PluginField("mha_type_id", np.array([get_mha_dtype(config)], np.int32), trt.PluginFieldType.INT32) |
| 381 | |
| 382 | pfc = trt.PluginFieldCollection([wbeta, wgamma, wwordemb, wtokemb, wposemb, output_fp16, mha_type]) |
| 383 | fn = emln_plg_creator.create_plugin("embeddings", pfc) |
| 384 | |
| 385 | input_ids = network.add_shuffle(input_ids) |
| 386 | input_ids.second_transpose = (1, 0) |
| 387 | segment_ids = network.add_shuffle(segment_ids) |
| 388 | segment_ids.second_transpose = (1, 0) |
| 389 | input_mask = network.add_shuffle(input_mask) |
| 390 | input_mask.second_transpose = (1, 0) |
| 391 | inputs = [input_ids.get_output(0), |
| 392 | segment_ids.get_output(0), |
| 393 | input_mask.get_output(0)] |
| 394 | emb_layer = network.add_plugin_v2(inputs, fn) |
| 395 | |
| 396 | if config.use_qat: |
| 397 | set_output_range(emb_layer, 1, 1) |
| 398 | set_output_name(emb_layer, "embeddings_", "output") |
| 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): |
no test coverage detected