Create the bert model
(config, init_dict, network, input_tensor, residual, mask_idx, cu_seqlens, max_seqlen)
| 309 | |
| 310 | |
| 311 | def bert_model(config, init_dict, network, input_tensor, residual, mask_idx, cu_seqlens, max_seqlen): |
| 312 | """ |
| 313 | Create the bert model |
| 314 | """ |
| 315 | prev_input = input_tensor |
| 316 | for layer in range(0, config.num_hidden_layers): |
| 317 | ss = "l{}_".format(layer) |
| 318 | out_layer = transformer_layer_opt(ss, config, init_dict, network, prev_input, residual, mask_idx, cu_seqlens, max_seqlen) |
| 319 | prev_input = out_layer.get_output(0) |
| 320 | # Skip reading residual from final layer |
| 321 | if config.use_megatron and (layer != config.num_hidden_layers - 1): |
| 322 | residual = out_layer.get_output(1) |
| 323 | |
| 324 | if config.use_qat: |
| 325 | dr_out = init_dict["bert_encoder_final_input_quantizer_amax"] |
| 326 | set_output_range(out_layer, dr_out) |
| 327 | |
| 328 | squad_logits = squad_output("cls_", config, init_dict, network, prev_input) |
| 329 | squad_logits_out = squad_logits.get_output(0) |
| 330 | network.mark_output(squad_logits_out) |
| 331 | |
| 332 | |
| 333 | def squad_output(prefix, config, init_dict, network, input_tensor): |
no test coverage detected