MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / squad_output

Function squad_output

demo/BERT/builder.py:324–342  ·  view source on GitHub ↗

Create the squad output

(prefix, config, init_dict, network, input_tensor)

Source from the content-addressed store, hash-verified

322 return prev_input
323
324def squad_output(prefix, config, init_dict, network, input_tensor):
325 """
326 Create the squad output
327 """
328
329 idims = input_tensor.shape
330 assert len(idims) == 5
331 B, S, hidden_size, _, _ = idims
332
333 W_out = init_dict[prefix + SQD_W]
334 B_out = init_dict[prefix + SQD_B]
335
336 W = network.add_constant((1, hidden_size, 2), W_out)
337 dense = network.add_fully_connected(input_tensor, 2, W_out, B_out)
338
339 OUT = network.add_shuffle(dense.get_output(0))
340 OUT.second_transpose = (1, 0, 2, 3, 4)
341 set_output_name(OUT, prefix, "squad_logits")
342 return OUT
343
344def 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.

Callers 1

build_engineFunction · 0.70

Calls 1

set_output_nameFunction · 0.70

Tested by

no test coverage detected