def test_trainable_varaible(self):
path = test_util.get_test_data_path("hub_module_v1_mini_train")
layer = hub_loader.HubKerasLayerV1V2(path, trainable=True)
self.assertLen(layer.trainable_variables, 2)
self.assertLen(layer.variables, 4)
layer = hub_loader.HubKerasLayerV1V2(path, trainable=False)
self.assertEmpty(layer.trainable_variables)
self.assertLen(layer.variables, 2)
def _create_model(self, hparams=None):
"""Creates the classifier model for retraining."""
hparams = self._get_hparams_or_default(hparams)
module_layer = hub_loader.HubKerasLayerV1V2(
self.model_spec.uri, trainable=hparams.do_fine_tuning)
return lib.build_model(module_layer, hparams,
self.model_spec.input_image_shape, self.num_classes)
def create_qa_model(bert_config,
max_seq_length,
initializer=None,
hub_module_url=None,
hub_module_trainable=True,
is_tf2=True):
"""Returns BERT qa model along with core BERT model to import weights.
Args:
bert_config: BertConfig, the config defines the core Bert model.
max_seq_length: integer, the maximum input sequence length.
initializer: Initializer for the final dense layer in the span labeler.
Defaulted to TruncatedNormal initializer.
hub_module_url: TF-Hub path/url to Bert module.
hub_module_trainable: True to finetune layers in the hub module.
is_tf2: boolean, whether the hub module is in TensorFlow 2.x format.
Returns:
A tuple of (1) keras model that outputs start logits and end logits and
(2) the core BERT transformer encoder.
"""
if initializer is None:
initializer = tf.keras.initializers.TruncatedNormal(
stddev=bert_config.initializer_range)
input_word_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_word_ids')
input_mask = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_mask')
input_type_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_type_ids')
if is_tf2:
core_model = hub.KerasLayer(hub_module_url, trainable=hub_module_trainable)
pooled_output, sequence_output = core_model(
[input_word_ids, input_mask, input_type_ids])
else:
bert_model = hub_loader.HubKerasLayerV1V2(
hub_module_url,
signature='tokens',
signature_outputs_as_dict=True,
trainable=hub_module_trainable)
outputs = bert_model({
'input_ids': input_word_ids,
'input_mask': input_mask,
'segment_ids': input_type_ids
})
pooled_output = outputs['pooled_output']
sequence_output = outputs['sequence_output']
bert_encoder = tf.keras.Model(
inputs=[input_word_ids, input_mask, input_type_ids],
outputs=[sequence_output, pooled_output],
name='core_model')
return models.BertSpanLabeler(
network=bert_encoder, initializer=initializer), bert_encoder
def test_trainable_varaible(self):
path = test_util.get_test_data_path("hub_module_v1_mini_train")
layer = hub_loader.HubKerasLayerV1V2(path, trainable=True)
# Checks trainable variables.
self.assertLen(layer.trainable_variables, 2)
self.assertEqual(layer.trainable_variables[0].name, "a:0")
self.assertEqual(layer.trainable_variables[1].name, "b:0")
self.assertEqual(layer.variables, layer.trainable_variables)
# Checks non-trainable variables.
self.assertEmpty(layer.non_trainable_variables)
layer = hub_loader.HubKerasLayerV1V2(path, trainable=False)
# Checks trainable variables.
self.assertEmpty(layer.trainable_variables)
# Checks non-trainable variables.
self.assertLen(layer.non_trainable_variables, 2)
self.assertEqual(layer.non_trainable_variables[0].name, "a:0")
self.assertEqual(layer.non_trainable_variables[1].name, "b:0")
self.assertEqual(layer.variables, layer.non_trainable_variables)
def create_model(self, hparams=None, with_loss_and_metrics=False):
"""Creates the classifier model for retraining."""
hparams = self._get_hparams_or_default(hparams)
module_layer = hub_loader.HubKerasLayerV1V2(
self.model_spec.uri, trainable=hparams.do_fine_tuning)
self.model = hub_lib.build_model(module_layer, hparams,
self.model_spec.input_image_shape,
self.num_classes)
if with_loss_and_metrics:
# Adds loss and metrics in the keras model.
self.model.compile(loss=tf.keras.losses.CategoricalCrossentropy(
label_smoothing=0.1),
metrics=['accuracy'])
def create_qa_model_from_squad(max_seq_length,
hub_module_url,
hub_module_trainable=True,
is_tf2=False):
"""Creates QA model the initialized from the model retrained on Squad dataset.
Args:
max_seq_length: integer, the maximum input sequence length.
hub_module_url: TF-Hub path/url to Bert module that's retrained on Squad
dataset.
hub_module_trainable: True to finetune layers in the hub module.
is_tf2: boolean, whether the hub module is in TensorFlow 2.x format.
Returns:
Keras model that outputs start logits and end logits.
"""
if is_tf2:
raise ValueError('Only supports to load TensorFlow 1.x hub module.')
input_word_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_word_ids')
input_mask = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_mask')
input_type_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_type_ids')
squad_bert = hub_loader.HubKerasLayerV1V2(
hub_module_url,
signature='squad',
signature_outputs_as_dict=True,
trainable=hub_module_trainable)
outputs = squad_bert({
'input_ids': input_word_ids,
'input_mask': input_mask,
'segment_ids': input_type_ids
})
start_logits = tf.keras.layers.Lambda(
tf.identity, name='start_positions')(
outputs['start_logits'])
end_logits = tf.keras.layers.Lambda(
tf.identity, name='end_positions')(
outputs['end_logits'])
return tf.keras.Model(
inputs=[input_word_ids, input_mask, input_type_ids],
outputs=[start_logits, end_logits])
def create_classifier_model(bert_config,
num_labels,
max_seq_length,
initializer=None,
hub_module_url=None,
hub_module_trainable=True,
is_tf2=True):
"""BERT classifier model in functional API style.
Construct a Keras model for predicting `num_labels` outputs from an input with
maximum sequence length `max_seq_length`.
Args:
bert_config: BertConfig, the config defines the core Bert model.
num_labels: integer, the number of classes.
max_seq_length: integer, the maximum input sequence length.
initializer: Initializer for the final dense layer in the span labeler.
Defaulted to TruncatedNormal initializer.
hub_module_url: TF-Hub path/url to Bert module.
hub_module_trainable: True to finetune layers in the hub module.
is_tf2: boolean, whether the hub module is in TensorFlow 2.x format.
Returns:
Combined prediction model (words, mask, type) -> (one-hot labels)
BERT sub-model (words, mask, type) -> (bert_outputs)
"""
if initializer is None:
initializer = tf.keras.initializers.TruncatedNormal(
stddev=bert_config.initializer_range)
input_word_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_word_ids')
input_mask = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_mask')
input_type_ids = tf.keras.layers.Input(
shape=(max_seq_length,), dtype=tf.int32, name='input_type_ids')
if is_tf2:
bert_model = hub.KerasLayer(hub_module_url, trainable=hub_module_trainable)
pooled_output, _ = bert_model([input_word_ids, input_mask, input_type_ids])
else:
bert_model = hub_loader.HubKerasLayerV1V2(
hub_module_url,
signature='tokens',
output_key='pooled_output',
trainable=hub_module_trainable)
pooled_output = bert_model({
'input_ids': input_word_ids,
'input_mask': input_mask,
'segment_ids': input_type_ids
})
output = tf.keras.layers.Dropout(rate=bert_config.hidden_dropout_prob)(
pooled_output)
output = tf.keras.layers.Dense(
num_labels,
kernel_initializer=initializer,
name='output',
activation='softmax',
dtype=tf.float32)(
output)
return tf.keras.Model(
inputs=[input_word_ids, input_mask, input_type_ids],
outputs=output), bert_model
def test_load_with_defaults(self, module_name, trainable): inputs, expected_outputs = 10., 11. # Test modules perform increment op. path = test_util.get_test_data_path(module_name) layer = hub_loader.HubKerasLayerV1V2(path, trainable=trainable) output = layer(inputs) self.assertEqual(output, expected_outputs)