Create the squad output
(prefix, config, init_dict, network, input_tensor)
| 331 | |
| 332 | |
| 333 | def squad_output(prefix, config, init_dict, network, input_tensor): |
| 334 | """ |
| 335 | Create the squad output |
| 336 | """ |
| 337 | hidden_size = config.hidden_size |
| 338 | |
| 339 | W_out = init_dict[prefix + SQD_W] |
| 340 | B_out = init_dict[prefix + SQD_B] |
| 341 | |
| 342 | if config.use_int8: |
| 343 | dense = network.add_convolution_nd(input_tensor, 2, (1, 1), W_out, B_out) |
| 344 | else: |
| 345 | dense = network.add_fully_connected(input_tensor, 2, W_out, B_out) |
| 346 | |
| 347 | OUT = network.add_shuffle(dense.get_output(0)) |
| 348 | if config.use_int8 and config.interleaved: |
| 349 | OUT.second_transpose = (1, 2, 0, 3) |
| 350 | else: |
| 351 | OUT.second_transpose = (1, 0, 2, 3) |
| 352 | set_output_name(OUT, prefix, "squad_logits") |
| 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,)) |
no test coverage detected