def smoke_test_model(model_path): try: resolved_model = hub.resolve(model_path) loader_impl.parse_saved_model(resolved_model) except Exception as e: # pylint: disable=broad-except return False, e return True, None
def build_encoding(encoding_strat, text_input): if encoding_strat == 'tfidf': encoder = tf.keras.layers.experimental.preprocessing.TextVectorization( max_tokens=VOCAB_SIZE, output_mode='tf-idf') encoder.adapt(train_dataset.map(lambda text, label: text)) net = encoder(text_input) return net elif encoding_strat == 'learned': encoder = tf.keras.layers.experimental.preprocessing.TextVectorization( max_tokens=VOCAB_SIZE, output_mode='int') encoder.adapt(train_dataset.map(lambda text, label: text)) net = encoder(text_input) net = tf.keras.layers.Embedding(input_dim=len( encoder.get_vocabulary()), output_dim=64, mask_zero=True)(net) net = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(10, name='LSTM'), name='Bidirectional')(net) return net elif encoding_strat == 'bert': load_options = tf.saved_model.LoadOptions( experimental_io_device='/job:localhost') preprocessor = hub.load( "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3", options=load_options) print("\n\n\nHub resolve:") print( hub.resolve( "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")) print("\n\n") # Step 1: tokenize batches of text inputs. tokenize = hub.KerasLayer(preprocessor.tokenize, name='tokenizer') tokenized_inputs = [ tokenize(text_input), ] # Step 2 (optional): modify tokenized inputs. # Step 3: pack input sequences for the Transformer encoder. seq_length = 250 # We filter out anything more earlier bert_pack_inputs = hub.KerasLayer( preprocessor.bert_pack_inputs, arguments=dict(seq_length=seq_length), name='input_packer') # Optional argument. encoder_inputs = bert_pack_inputs(tokenized_inputs) encoder = hub.KerasLayer( "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-6_H-512_A-8/1", trainable=False, name='BERT_encoder') outputs = encoder(encoder_inputs) sequence_output = outputs["sequence_output"] net = tf.keras.layers.Bidirectional( tf.keras.layers.LSTM(10, name='LSTM'), name='Bidirectional')(sequence_output) return net else: raise ValueError(f"Invalid encoding strategy - given {encoding_strat}")
def __init__(self, num_units, CLS, SEP, PAD, UNK, min_len, max_len, bert_hub_url, **kwargs): """ Initializes the layer For more details on parameters, check the argument parser in run_detext.py """ super(BertLayer, self).__init__() self.text_ftr_size = num_units # Load pretrained model using hub with tf.name_scope(BERT_VAR_PREFIX): self.text_encoder = hub.KerasLayer(hub.resolve(bert_hub_url), trainable=True) self.last_layer_cls = tf.keras.layers.Dense(num_units, activation='tanh', use_bias=True, name=BERT_VAR_PREFIX + '_last_layer_cls') self.preprocess_layer = BertPreprocessLayer(self.text_encoder, max_len, min_len, CLS, SEP, PAD, UNK) self._pad_id = self.preprocess_layer.pad_id() self._pad_id = tf.cast(self._pad_id, tf.int32) self.dropout = tf.keras.layers.Dropout(0.1)
def convert_tf_hub_module(module_handle, output_dir, signature='default', saved_model_tags='serve', quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, control_flow_v2=False, experiments=False, metadata=None): """Conversion for TF Hub modules V1 and V2. See convert_tf_hub_module and convert_tf_saved_model. Args: module_path: string Path to the module. output_dir: string The name of the output directory. The directory will consist of - a file named 'model.json' - possibly sharded binary weight files. signature: string Signature to load. saved_model_tags: tags of the GraphDef to load. Defaults to ''. quantization_dtype_map: A mapping from dtype (`uint8`, `uint16`, `float16`) to weights names. The weight mapping supports wildcard substitution. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to strip debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. experiments: Bool enable experimental features. metadata: User defined metadata map. """ module_path = hub.resolve(module_handle) # TODO(vbardiovskyg): We can remove this v1 code path once loading of all v1 # modules is fixed on the TF side, or once the modules we cannot load become # replaced with newer versions. if tf.io.gfile.exists(os.path.join(module_path, _HUB_V1_MODULE_PB)): print("Loading the module using TF 1.X interface from %s." % module_path) convert_tf_hub_module_v1(module_path, output_dir, signature, quantization_dtype_map, skip_op_check, strip_debug_ops, weight_shard_size_bytes, experiments=experiments, metadata=metadata) else: print("Loading the module using TF 2.X interface from %s." % module_path) if signature is None: signature = 'default' convert_tf_saved_model(saved_model_dir=module_path, output_dir=output_dir, signature_def=signature, saved_model_tags=saved_model_tags, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, control_flow_v2=control_flow_v2, experiments=experiments, metadata=metadata)
def test_resolve(self): with tf.Graph().as_default(): self._generate_module() module_dir = hub.resolve( "http://localhost:%d/test_module.tgz" % self.server_port) self.assertIn(tempfile.gettempdir(), module_dir) module_files = sorted(tf_v1.gfile.ListDirectory(module_dir)) self.assertEqual( ["assets", "saved_model.pb", "tfhub_module.pb", "variables"], module_files)
def assert_can_resolve_asset(self, asset_path: str): """Attempt to hub.resolve the given asset path.""" try: resolved_model = hub.resolve(asset_path) loader_impl.parse_saved_model(resolved_model) _validate_file_paths(resolved_model) except Exception as e: # pylint: disable=broad-except raise MarkdownDocumentationError( f"The model on path {asset_path} failed to parse. Please make sure " "that the asset-path metadata points to a valid TF2 SavedModel or a " "TF1 Hub module, compressed as described in section 'Model' of " f"README.md. Underlying reason for failure: {e}.")
def __init__(self, tfhub_handle = DEFAULT_TOKENIZER_TFHUB_HANDLE): preprocessor = hub.load(tfhub_handle) model_path = hub.resolve(tfhub_handle) vocab_file_path = os.path.join(model_path, 'assets/vocab.txt') with tf.io.gfile.GFile(vocab_file_path, 'r') as vocab_file: self.vocab = [token.strip() for token in vocab_file.readlines()] text_input = tf.keras.layers.Input(shape=(), dtype=tf.string) tokenize_layer = hub.KerasLayer(preprocessor.tokenize) outputs = tokenize_layer(text_input) self.model = tf.keras.Model( inputs=text_input, outputs=outputs, name='tokenizer')
def assert_can_resolve_asset(self, asset_path: str): """Attempt to hub.resolve the given asset path.""" try: resolved_model = hub.resolve(asset_path) loader_impl.parse_saved_model(resolved_model) _validate_file_paths(resolved_model) except Exception as e: # pylint: disable=broad-except raise MarkdownDocumentationError( f"The model on path {asset_path} failed to parse. Please make sure " "that the asset-path metadata points to a valid TF2 SavedModel or a " "TF1 Hub module as described on " "https://www.tensorflow.org/hub/exporting_tf2_saved_model. " f"Underlying reason for failure: {e}.")
def convert_tf_hub_module(module_handle, output_dir, signature='default', saved_model_tags='serve', quantization_dtype=None, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4): """Conversion for TF Hub modules V1 and V2. See convert_tf_hub_module and convert_tf_saved_model. Args: module_path: string Path to the module. output_dir: string The name of the output directory. The directory will consist of - a file named 'model.json' - possibly sharded binary weight files. signature: string Signature to load. saved_model_tags: tags of the GraphDef to load. Defaults to ''. skip_op_check: Bool whether to skip the op check. strip_debug_ops: Bool whether to strip debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. """ module_path = hub.resolve(module_handle) # TODO(vbardiovskyg): We can remove this v1 code path once loading of all v1 # modules is fixed on the TF side, or once the modules we cannot load become # replaced with newer versions. if tf.io.gfile.exists(os.path.join(module_path, _HUB_V1_MODULE_PB)): print("Loading the module using TF 1.X interface from %s." % module_path) convert_tf_hub_module_v1(module_path, output_dir, signature, quantization_dtype, skip_op_check, strip_debug_ops, weight_shard_size_bytes) else: print("Loading the module using TF 2.X interface from %s." % module_path) if signature is None: signature = 'default' convert_tf_saved_model(saved_model_dir=module_path, output_dir=output_dir, signature_def=signature, saved_model_tags=saved_model_tags, quantization_dtype=quantization_dtype, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes)
def import_from_tfhub( identifier: t.Union[str, "HubModule"], name: t.Optional[str] = None, labels: t.Optional[t.Dict[str, str]] = None, custom_objects: t.Optional[t.Dict[str, t.Any]] = None, metadata: t.Optional[t.Dict[str, t.Any]] = None, model_store: "ModelStore" = Provide[BentoMLContainer.model_store], ) -> Tag: """ Import a model from `Tensorflow Hub <https://tfhub.dev/>`_ to BentoML modelstore. Args: identifier (:code:`Union[str, tensorflow_hub.Module, tensorflow_hub.KerasLayer]`): Identifier accepts two type of inputs: - if `type` of :code:`identifier` either of type :code:`tensorflow_hub.Module` (**legacy** `tensorflow_hub`) or :code:`tensorflow_hub.KerasLayer` (`tensorflow_hub`), then we will save the given model to a :code:`SavedModel` format. - if `type` of :code:`identifier` is a :obj:`str`, we assume that this is the URI retrieved from Tensorflow Hub. We then clean the given URI, and get a local copy of a given model to BentoML modelstore. name (:code:`str`, `optional`, defaults to `None`): An optional name for the model. If :code:`identifier` is a :obj:`str`, then name can be autogenerated from the given URI. name (:code:`str`, `optional`, default to `None`): Optional name for the saved model. If None, then name will be generated from :code:`identifier`. labels (:code:`Dict[str, str]`, `optional`, default to :code:`None`): user-defined labels for managing models, e.g. team=nlp, stage=dev custom_objects (:code:`Dict[str, Any]]`, `optional`, default to :code:`None`): user-defined additional python objects to be saved alongside the model, e.g. a tokenizer instance, preprocessor function, model configuration json metadata (:code:`Dict[str, Any]`, `optional`, default to :code:`None`): Custom metadata for given model. model_store (:mod:`~bentoml._internal.models.store.ModelStore`, default to :mod:`BentoMLContainer.model_store`): BentoML modelstore, provided by DI Container. Returns: :obj:`~bentoml.Tag`: A :obj:`~bentoml.Tag` object that can be used to retrieve the model with :func:`bentoml.tensorflow.load`: Example for importing a model from Tensorflow Hub: .. code-block:: python import tensorflow_text as text # noqa # pylint: disable import bentoml tag = bentoml.tensorflow_v1.import_from_tfhub("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3") # load model back with `load`: model = bentoml.tensorflow_v1.load(tag, load_as_hub_module=True) Example for importing a custom Tensorflow Hub model: .. code-block:: python import tensorflow as tf import tensorflow_hub as hub import bentoml # Simple toy Tensorflow Hub model def _plus_one_model_tf1(): def plus_one(): x = tf.placeholder(dtype=tf.float32, name="x") y = x + 1 hub.add_signature(inputs=x, outputs=y) spec = hub.create_module_spec(plus_one) with tf.get_default_graph().as_default(): module = hub.Module(spec, trainable=True) return module model = _plus_one_model_tf1() # retrieve the given tag: tag = bentoml.tensorflow_v1.import_from_tfhub(model) """ # noqa if hub is None: raise MissingDependencyException("""\ `tensorflow_hub` does not exists. Make sure to `pip install --upgrade tensorflow_hub` before using. """) context: t.Dict[str, t.Any] = { "framework_name": "tensorflow", "pip_dependencies": [ f"tensorflow=={get_tf_version()}", f"tensorflow_hub=={importlib_metadata.version('tensorflow_hub')}", ], "import_from_tfhub": True, } if name is None: if isinstance(identifier, str): name = _clean_name(identifier) else: name = f"{identifier.__class__.__name__}_{uuid.uuid4().hex[:5].upper()}" with bentoml.models.create( name, module=MODULE_NAME, options=None, context=context, labels=labels, custom_objects=custom_objects, metadata=metadata, ) as _model: if isinstance(identifier, str): current_cache_dir = os.environ.get("TFHUB_CACHE_DIR") os.environ["TFHUB_CACHE_DIR"] = _model.path fpath: str = resolve(identifier) folder = fpath.split("/")[-1] _model.info.options = {"model": identifier, "local_path": folder} if current_cache_dir is not None: os.environ["TFHUB_CACHE_DIR"] = current_cache_dir else: if hasattr(identifier, "export"): # hub.Module.export() with tf.Session( graph=tf.get_default_graph()) as sess: # type: ignore sess.run(tf.global_variables_initializer()) # type: ignore identifier.export(_model.path, sess) # type: ignore else: tf.saved_model.save(identifier, _model.path) _model.info.options = { "model": identifier.__class__.__name__, "local_path": ".", } return _model.tag
class TestBertLayer(tf.test.TestCase, DataSetup): """Unit test for bert_layer.py""" # Bert setup bert_hub_layer = hub.KerasLayer(hub.resolve(DataSetup.bert_hub_url), trainable=True) bert_vocab_file = bert_hub_layer.resolved_object.vocab_file.asset_path.numpy( ).decode("utf-8") bert_vocab_table = read_vocab(bert_vocab_file) bert_PAD_ID = bert_vocab_table[DataSetup.PAD] bert_SEP_ID = bert_vocab_table[DataSetup.SEP] bert_CLS_ID = bert_vocab_table[DataSetup.CLS] bert_UNK_ID = bert_vocab_table[DataSetup.UNK] # SentencePiece setup sentencepiece_hub_layer = hub.KerasLayer( hub.resolve(DataSetup.libert_sp_hub_url)) tokenizer_file = sentencepiece_hub_layer.resolved_object.tokenizer_file.asset_path.numpy( ).decode("utf-8") with tf.io.gfile.GFile(tokenizer_file, 'rb') as f_handler: sp_model = f_handler.read() sentencepiece_tokenizer = tf_text.SentencepieceTokenizer(model=sp_model, out_type=tf.int32) sentencepiece_vocab_tf_table = create_tf_vocab_from_sp_tokenizer( sp_tokenizer=sentencepiece_tokenizer, num_oov_buckets=1) sentencepiece_PAD_ID = sentencepiece_vocab_tf_table.lookup( tf.constant(DataSetup.PAD)).numpy() sentencepiece_SEP_ID = sentencepiece_vocab_tf_table.lookup( tf.constant(DataSetup.SEP)).numpy() sentencepiece_CLS_ID = sentencepiece_vocab_tf_table.lookup( tf.constant(DataSetup.CLS)).numpy() sentencepiece_UNK_ID = sentencepiece_vocab_tf_table.lookup( tf.constant(DataSetup.UNK)).numpy() # Space setup space_hub_layer = hub.KerasLayer( hub.resolve(DataSetup.libert_space_hub_url)) tokenizer_file = space_hub_layer.resolved_object.tokenizer_file.asset_path.numpy( ).decode("utf-8") space_vocab = read_vocab(tokenizer_file) space_PAD_ID = space_vocab[DataSetup.PAD] space_SEP_ID = space_vocab[DataSetup.SEP] space_CLS_ID = space_vocab[DataSetup.CLS] space_UNK_ID = space_vocab[DataSetup.UNK] # Hyperparameters setup num_units = 16 pad_id = 0 num_doc_fields = 2 min_len = 3 max_len = 8 layer = bert_layer.BertLayer(num_units, DataSetup.CLS, DataSetup.SEP, DataSetup.PAD, DataSetup.UNK, min_len, max_len, DataSetup.bert_hub_url) def testBertLayer(self): """Test Bert layer """ for hub_url in [self.bert_hub_url]: self._testBertLayer(hub_url) def _testBertLayer(self, hub_url): query = self.query doc_fields = [self.ranking_doc_field1, self.ranking_doc_field2] user_fields = [query, query, query] query_ftrs, doc_ftrs, user_ftrs = self.layer( { InputFtrType.QUERY_COLUMN_NAME: query, InputFtrType.DOC_TEXT_COLUMN_NAMES: doc_fields, InputFtrType.USER_TEXT_COLUMN_NAMES: user_fields }, False) text_ftr_size = self.num_units self.assertEqual(text_ftr_size, self.layer.text_ftr_size) self.assertAllEqual(query_ftrs.shape, [2, self.layer.text_ftr_size]) self.assertAllEqual(doc_ftrs.shape, [2, 3, 2, self.layer.text_ftr_size]) self.assertAllEqual(user_ftrs.shape, [2, 3, self.layer.text_ftr_size]) # 1st query, 2nd doc, 2nd field should be the same as 2nd query, 1st doc, 2nd field self.assertAllEqual(doc_ftrs[0, 1, 1], doc_ftrs[1, 0, 1]) # 1st query, 1st doc, 1st field should be the same as 1st query, 1st doc, 2nd field self.assertAllEqual(doc_ftrs[0, 0, 0], doc_ftrs[0, 0, 1]) # 1st query, 1st doc, 2st field should NOT be the same as 1st query, 2st doc, 2nd field self.assertNotAllClose(doc_ftrs[0, 1, 0], doc_ftrs[0, 1, 1]) def testGetInputIds(self): """Tests get_input_ids() """ query = [[1, 2, 3], [2, 4, 3]] doc_field1 = [[[1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0]], [[2, 4, 3, 1], [1, 3, 3, 1], [1, 3, 3, 1]]] doc_field2 = [[[1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0]], [[20, 5, 3, 1], [5, 6, 1, 1], [5, 6, 1, 1]]] query = tf.constant(query, dtype=tf.int32) doc_field1 = tf.constant(doc_field1, dtype=tf.int32) doc_field2 = tf.constant(doc_field2, dtype=tf.int32) doc_fields = [doc_field1, doc_field2] user_fields = None max_text_len, max_text_len_array = bert_layer.BertLayer.get_input_max_len( query, doc_fields, user_fields) bert_input_ids = bert_layer.BertLayer.get_bert_input_ids( query, doc_fields, user_fields, self.pad_id, max_text_len, max_text_len_array) # Check bert input ids self.assertAllEqual( bert_input_ids, [[1, 2, 3, 0], [2, 4, 3, 0], [1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0], [2, 4, 3, 1], [1, 3, 3, 1], [1, 3, 3, 1], [1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0], [20, 5, 3, 1], [5, 6, 1, 1], [5, 6, 1, 1]]) def testPrerocessQuery(self): """Tests _preprocess_query function of bert layer""" query = tf.constant(['batch 1 user 1 build', 'batch 2 user 2 word'], dtype=tf.string) expected = tf.constant( [[ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 4, 2 ], [ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 5, 2 ]], dtype=tf.int32) self.assertAllEqual(expected, self.layer._preprocess_query(query)) def testPrerocessUsr(self): """Tests _preprocess_user function of bert layer""" user_fields = [ tf.constant(['batch 1 user 1 build', 'batch 2 user 2 word'], dtype=tf.string) ] expected = [ tf.constant([[ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 4, 2 ], [ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 5, 2 ]], dtype=tf.int32) ] self.assertAllEqual(expected, self.layer._preprocess_user(user_fields)) def testPrerocessDoc(self): """Tests _preprocess_doc function of bert layer""" doc_fields = [ tf.constant([['batch 1 doc 1 build', 'batch 1 doc 2'], ['batch 2 doc 1 word', 'batch 2 doc 2']], dtype=tf.string) ] expected = [ tf.constant([[[ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 4, 2 ], [ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 2, 3 ]], [[ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 5, 2 ], [ self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, self.bert_UNK_ID, 2, 3 ]]], dtype=tf.int32) ] self.assertAllEqual(expected, self.layer._preprocess_doc(doc_fields)) def testBertPreprocessLayerWordPiece(self): """Tests BertPreprocessLayer with wordpiece tokenizer""" preprocess_layer = bert_layer.BertPreprocessLayer( self.bert_hub_layer, self.max_len, self.min_len, self.CLS, self.SEP, self.PAD, self.UNK) sentences = tf.constant( ['test sent1', 'build build build build sent2']) expected = tf.constant([[ self.bert_CLS_ID, 8, self.bert_UNK_ID, self.bert_SEP_ID, self.bert_PAD_ID, self.bert_PAD_ID, self.bert_PAD_ID ], [self.bert_CLS_ID, 4, 4, 4, 4, self.bert_UNK_ID, self.bert_SEP_ID]], dtype=tf.int32) outputs = preprocess_layer(sentences) self.assertAllEqual(expected, outputs) def testBertPreprocessLayerSentencePiece(self): """Tests BertPreprocessLayer with sentencepiece tokenizer""" preprocess_layer = bert_layer.BertPreprocessLayer( self.sentencepiece_hub_layer, self.max_len, self.min_len, self.CLS, self.SEP, self.PAD, self.UNK) sentences = tf.constant( ['TEST sent1', 'build build build build sent2']) expected = tf.constant([[ self.sentencepiece_CLS_ID, 557, 4120, 29900, self.sentencepiece_SEP_ID, self.sentencepiece_PAD_ID, self.sentencepiece_PAD_ID, self.sentencepiece_PAD_ID ], [ self.sentencepiece_CLS_ID, 671, 671, 671, 671, 4120, 29904, self.sentencepiece_SEP_ID ]], dtype=tf.int32) outputs = preprocess_layer(sentences) self.assertAllEqual(expected, outputs) def testBertPreprocessLayerSpace(self): """Tests BertPreprocessLayer with space tokenizer""" preprocess_layer = bert_layer.BertPreprocessLayer( self.space_hub_layer, self.max_len, self.min_len, self.CLS, self.SEP, self.PAD, self.UNK) sentences = tf.constant( ['test sent1', 'build build build build sent2']) expected = tf.constant([[ self.space_CLS_ID, 8, self.space_UNK_ID, self.space_SEP_ID, self.space_PAD_ID, self.space_PAD_ID, self.space_PAD_ID ], [ self.space_CLS_ID, 4, 4, 4, 4, self.space_UNK_ID, self.space_SEP_ID ]], dtype=tf.int32) outputs = preprocess_layer(sentences) self.assertAllEqual(expected, outputs) def testBertPreprocessLayerAdjustLen(self): """Tests adjust_len function of BertPreprocessLayer""" sentences = tf.constant( ['test sent1', 'build build build build sent2']) min_len = 12 max_len = 16 layer = bert_layer.BertPreprocessLayer(self.bert_hub_layer, max_len, min_len, self.CLS, self.SEP, self.PAD, self.UNK) outputs = layer(sentences) shape = tf.shape(outputs) self.assertAllEqual(shape, tf.constant([2, 12])) min_len = 0 max_len = 1 layer = bert_layer.BertPreprocessLayer(self.bert_hub_layer, max_len, min_len, self.CLS, self.SEP, self.PAD, self.UNK) outputs = layer(sentences) shape = tf.shape(outputs) self.assertAllEqual(shape, tf.constant([2, 1]))
def model_fn(features, mode, params): """A function for applying hub module that follows Estimator API.""" hub_module = params["hub_module"] finetune_layer = params["finetune_layer"] num_classes = params["num_classes"] initial_learning_rate = params["initial_learning_rate"] momentum = params["momentum"] lr_decay_factor = params["lr_decay_factor"] decay_steps = params["decay_steps"] warmup_steps = params["warmup_steps"] is_training = (mode == tf.estimator.ModeKeys.TRAIN) module_path = hub.resolve(hub_module) is_legacy_hub_module = tf.io.gfile.exists( os.path.join(module_path, "tfhub_module.pb")) if is_legacy_hub_module: module = hub.Module(hub_module, trainable=is_training, tags={"train"} if is_training else None) pre_logits = module(features["image"], signature=params["hub_module_signature"], as_dict=True)[finetune_layer] else: module = hub.load(hub_module) tf.get_collection_ref(tf.GraphKeys.TRAINABLE_VARIABLES).extend( module.trainable_variables) pre_logits = module(features["image"], training=is_training) num_dim_pre_logits = len(pre_logits.get_shape().as_list()) if num_dim_pre_logits == 4: pre_logits = tf.squeeze(pre_logits, [1, 2]) elif num_dim_pre_logits != 2: raise ValueError("Invalid number of dimensions in the representation " "layer: {}, but only 2 or 4 are allowed".format( num_dim_pre_logits)) logits = tf.layers.dense(pre_logits, units=num_classes, kernel_initializer=tf.zeros_initializer(), name="linear_head") loss = tf.nn.sparse_softmax_cross_entropy_with_logits( logits=logits, labels=features["label"]) loss = tf.reduce_mean(loss) def accuracy_metric(logits, labels): return {"accuracy": tf.metrics.accuracy( labels=labels, predictions=tf.argmax(logits, axis=-1))} eval_metrics = (accuracy_metric, [logits, features["label"]]) if mode == tf.estimator.ModeKeys.EVAL: if params["tpu_name"] is not None: return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=loss, eval_metrics=eval_metrics) else: return tf.estimator.EstimatorSpec( mode=mode, loss=loss, eval_metric_ops=eval_metrics[0](*eval_metrics[1])) elif mode == tf.estimator.ModeKeys.TRAIN: train_op = trainer.get_train_op(loss, initial_learning_rate, momentum, lr_decay_factor, decay_steps, warmup_steps, use_tpu=params["tpu_name"] is not None) spec_type = (tf.contrib.tpu.TPUEstimatorSpec if params["tpu_name"] is not None else tf.estimator.EstimatorSpec) return spec_type(mode=mode, loss=loss, train_op=train_op) else: raise ValueError("Unsupported mode %s" % mode)