コード例 #1
0
 def test_export_saved_model(self):
     model = Bert(
         vocab_size=21128, num_layers=12, num_attention_heads=8, return_states=True, return_attention_weights=True
     )
     input_ids, segment_ids, input_mask = model.dummy_inputs()
     model(inputs=[input_ids, segment_ids, input_mask])
     model.summary()
     model.save("models/export/2", include_optimizer=False)
コード例 #2
0
    def _check_bert_outputs(self, return_states=False, return_attention_weights=False):
        NUM_LAYERS = 4
        model = Bert(
            vocab_size=100,
            num_layers=NUM_LAYERS,
            return_states=return_states,
            return_attention_weights=return_attention_weights,
        )
        input_ids, segment_ids, attn_mask = self._build_bert_inputs()
        outputs = model(inputs=[input_ids, segment_ids, attn_mask])
        sequence_outputs, pooled_outputs = outputs[0], outputs[1]
        self.assertAllEqual([2, 16, 768], sequence_outputs.shape)
        self.assertAllEqual([2, 768], pooled_outputs.shape)

        all_states, all_attn_weights = None, None
        if return_states and return_attention_weights:
            self.assertEqual(4, len(outputs))
            all_states, all_attn_weights = outputs[2], outputs[3]
        elif return_states and not return_attention_weights:
            self.assertEqual(3, len(outputs))
            all_states = outputs[2]
        elif not return_states and return_attention_weights:
            self.assertEqual(3, len(outputs))
            all_attn_weights = outputs[2]
        else:
            self.assertEqual(2, len(outputs))

        if all_states is not None:
            # self.assertEqual(2, len(all_states))
            # for state in all_states:
            #     self.assertAllEqual([2, 16, 768], state.shape)
            self.assertAllEqual([2, NUM_LAYERS, 16, 768], all_states.shape)

        if all_attn_weights is not None:
            # self.assertEqual(2, len(all_attn_weights))
            # for attention in all_attn_weights:
            #     self.assertAllEqual([2, 8, 16, 16], attention.shape)
            self.assertAllEqual([2, NUM_LAYERS, 8, 16, 16], all_attn_weights.shape)
コード例 #3
0
 def test_build_model(self):
     model = Bert(vocab_size=21128)
     input_ids, segment_ids, input_mask = model.dummy_inputs()
     model(inputs=[input_ids, segment_ids, input_mask])
     model.summary()
コード例 #4
0
 def test_bert_config(self):
     model = Bert(vocab_size=100, num_layers=2, return_states=True, return_attention_weights=True)
     config = model.get_config()
     print(config)