diff --git a/mediapipe/model_maker/python/text/text_classifier/BUILD b/mediapipe/model_maker/python/text/text_classifier/BUILD index 26412d2c..64ace4ba 100644 --- a/mediapipe/model_maker/python/text/text_classifier/BUILD +++ b/mediapipe/model_maker/python/text/text_classifier/BUILD @@ -31,11 +31,11 @@ py_library( visibility = ["//visibility:public"], deps = [ ":dataset", + ":hyperparameters", ":model_options", ":model_spec", ":text_classifier", ":text_classifier_options", - "//mediapipe/model_maker/python/core:hyperparameters", ], ) @@ -45,12 +45,18 @@ py_library( deps = ["//mediapipe/model_maker/python/text/core:bert_model_options"], ) +py_library( + name = "hyperparameters", + srcs = ["hyperparameters.py"], + deps = ["//mediapipe/model_maker/python/core:hyperparameters"], +) + py_library( name = "model_spec", srcs = ["model_spec.py"], deps = [ + ":hyperparameters", ":model_options", - "//mediapipe/model_maker/python/core:hyperparameters", "//mediapipe/model_maker/python/core/utils:file_util", "//mediapipe/model_maker/python/text/core:bert_model_spec", ], @@ -61,9 +67,9 @@ py_test( srcs = ["model_spec_test.py"], tags = ["requires-net:external"], deps = [ + ":hyperparameters", ":model_options", ":model_spec", - "//mediapipe/model_maker/python/core:hyperparameters", ], ) @@ -100,9 +106,9 @@ py_library( name = "text_classifier_options", srcs = ["text_classifier_options.py"], deps = [ + ":hyperparameters", ":model_options", ":model_spec", - "//mediapipe/model_maker/python/core:hyperparameters", ], ) @@ -111,11 +117,11 @@ py_library( srcs = ["text_classifier.py"], deps = [ ":dataset", + ":hyperparameters", ":model_options", ":model_spec", ":preprocessor", ":text_classifier_options", - "//mediapipe/model_maker/python/core:hyperparameters", "//mediapipe/model_maker/python/core/data:dataset", "//mediapipe/model_maker/python/core/tasks:classifier", "//mediapipe/model_maker/python/core/utils:metrics", diff --git a/mediapipe/model_maker/python/text/text_classifier/__init__.py b/mediapipe/model_maker/python/text/text_classifier/__init__.py index 4df3a771..7eb0f925 100644 --- a/mediapipe/model_maker/python/text/text_classifier/__init__.py +++ b/mediapipe/model_maker/python/text/text_classifier/__init__.py @@ -13,19 +13,23 @@ # limitations under the License. """MediaPipe Public Python API for Text Classifier.""" -from mediapipe.model_maker.python.core import hyperparameters from mediapipe.model_maker.python.text.text_classifier import dataset +from mediapipe.model_maker.python.text.text_classifier import hyperparameters from mediapipe.model_maker.python.text.text_classifier import model_options from mediapipe.model_maker.python.text.text_classifier import model_spec from mediapipe.model_maker.python.text.text_classifier import text_classifier from mediapipe.model_maker.python.text.text_classifier import text_classifier_options -HParams = hyperparameters.BaseHParams + +AverageWordEmbeddingHParams = hyperparameters.AverageWordEmbeddingHParams +AverageWordEmbeddingModelOptions = ( + model_options.AverageWordEmbeddingModelOptions +) +BertOptimizer = hyperparameters.BertOptimizer +BertHParams = hyperparameters.BertHParams +BertModelOptions = model_options.BertModelOptions CSVParams = dataset.CSVParameters Dataset = dataset.Dataset -AverageWordEmbeddingModelOptions = ( - model_options.AverageWordEmbeddingModelOptions) -BertModelOptions = model_options.BertModelOptions SupportedModels = model_spec.SupportedModels TextClassifier = text_classifier.TextClassifier TextClassifierOptions = text_classifier_options.TextClassifierOptions diff --git a/mediapipe/model_maker/python/text/text_classifier/hyperparameters.py b/mediapipe/model_maker/python/text/text_classifier/hyperparameters.py new file mode 100644 index 00000000..ae0a9a62 --- /dev/null +++ b/mediapipe/model_maker/python/text/text_classifier/hyperparameters.py @@ -0,0 +1,54 @@ +# Copyright 2023 The MediaPipe Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Hyperparameters for training object detection models.""" + +import dataclasses +import enum +from typing import Union + +from mediapipe.model_maker.python.core import hyperparameters as hp + + +@dataclasses.dataclass +class AverageWordEmbeddingHParams(hp.BaseHParams): + """The hyperparameters for an AverageWordEmbeddingClassifier.""" + + +@enum.unique +class BertOptimizer(enum.Enum): + """Supported Optimizers for Bert Text Classifier.""" + + ADAMW = "adamw" + LAMB = "lamb" + + +@dataclasses.dataclass +class BertHParams(hp.BaseHParams): + """The hyperparameters for a Bert Classifier. + + Attributes: + learning_rate: Learning rate to use for gradient descent training. + batch_size: Batch size for training. + epochs: Number of training iterations over the dataset. + optimizer: Optimizer to use for training. Only supported values are "adamw" + and "lamb". + """ + + learning_rate: float = 3e-5 + batch_size: int = 48 + epochs: int = 2 + optimizer: BertOptimizer = BertOptimizer.ADAMW + + +HParams = Union[BertHParams, AverageWordEmbeddingHParams] diff --git a/mediapipe/model_maker/python/text/text_classifier/model_spec.py b/mediapipe/model_maker/python/text/text_classifier/model_spec.py index e947f8c1..452e2267 100644 --- a/mediapipe/model_maker/python/text/text_classifier/model_spec.py +++ b/mediapipe/model_maker/python/text/text_classifier/model_spec.py @@ -17,13 +17,11 @@ import dataclasses import enum import functools -from mediapipe.model_maker.python.core import hyperparameters as hp from mediapipe.model_maker.python.core.utils import file_util from mediapipe.model_maker.python.text.core import bert_model_spec +from mediapipe.model_maker.python.text.text_classifier import hyperparameters as hp from mediapipe.model_maker.python.text.text_classifier import model_options as mo -# BERT-based text classifier spec inherited from BertModelSpec -BertClassifierSpec = bert_model_spec.BertModelSpec MOBILEBERT_TINY_FILES = file_util.DownloadedFiles( 'text_classifier/mobilebert_tiny', @@ -31,6 +29,12 @@ MOBILEBERT_TINY_FILES = file_util.DownloadedFiles( is_folder=True, ) +EXBERT_FILES = file_util.DownloadedFiles( + 'text_classifier/exbert', + 'https://storage.googleapis.com/mediapipe-assets/exbert.tar.gz', + is_folder=True, +) + @dataclasses.dataclass class AverageWordEmbeddingClassifierSpec: @@ -43,27 +47,53 @@ class AverageWordEmbeddingClassifierSpec: """ # `learning_rate` is unused for the average word embedding model - hparams: hp.BaseHParams = hp.BaseHParams( - epochs=10, batch_size=32, learning_rate=0) + hparams: hp.AverageWordEmbeddingHParams = hp.AverageWordEmbeddingHParams( + epochs=10, batch_size=32, learning_rate=0 + ) model_options: mo.AverageWordEmbeddingModelOptions = ( mo.AverageWordEmbeddingModelOptions()) name: str = 'AverageWordEmbedding' - average_word_embedding_classifier_spec = functools.partial( AverageWordEmbeddingClassifierSpec) + +@dataclasses.dataclass +class BertClassifierSpec(bert_model_spec.BertModelSpec): + """Specification for a Bert classifier model. + + Only overrides the hparams attribute since the rest of the attributes are + inherited from the BertModelSpec. + """ + + hparams: hp.BertHParams = hp.BertHParams() + + mobilebert_classifier_spec = functools.partial( BertClassifierSpec, downloaded_files=MOBILEBERT_TINY_FILES, - hparams=hp.BaseHParams( + hparams=hp.BertHParams( epochs=3, batch_size=48, learning_rate=3e-5, distribution_strategy='off' ), name='MobileBert', tflite_input_name={ 'ids': 'serving_default_input_1:0', - 'mask': 'serving_default_input_3:0', 'segment_ids': 'serving_default_input_2:0', + 'mask': 'serving_default_input_3:0', + }, +) + +exbert_classifier_spec = functools.partial( + BertClassifierSpec, + downloaded_files=EXBERT_FILES, + hparams=hp.BertHParams( + epochs=3, batch_size=48, learning_rate=3e-5, distribution_strategy='off' + ), + name='ExBert', + tflite_input_name={ + 'ids': 'serving_default_input_1:0', + 'segment_ids': 'serving_default_input_2:0', + 'mask': 'serving_default_input_3:0', }, ) @@ -73,3 +103,4 @@ class SupportedModels(enum.Enum): """Predefined text classifier model specs supported by Model Maker.""" AVERAGE_WORD_EMBEDDING_CLASSIFIER = average_word_embedding_classifier_spec MOBILEBERT_CLASSIFIER = mobilebert_classifier_spec + EXBERT_CLASSIFIER = exbert_classifier_spec diff --git a/mediapipe/model_maker/python/text/text_classifier/model_spec_test.py b/mediapipe/model_maker/python/text/text_classifier/model_spec_test.py index a8d40558..7c45a267 100644 --- a/mediapipe/model_maker/python/text/text_classifier/model_spec_test.py +++ b/mediapipe/model_maker/python/text/text_classifier/model_spec_test.py @@ -19,7 +19,7 @@ from unittest import mock as unittest_mock import tensorflow as tf -from mediapipe.model_maker.python.core import hyperparameters as hp +from mediapipe.model_maker.python.text.text_classifier import hyperparameters as hp from mediapipe.model_maker.python.text.text_classifier import model_options as classifier_model_options from mediapipe.model_maker.python.text.text_classifier import model_spec as ms @@ -57,11 +57,13 @@ class ModelSpecTest(tf.test.TestCase): seq_len=128, do_fine_tuning=True, dropout_rate=0.1)) self.assertEqual( model_spec_obj.hparams, - hp.BaseHParams( + hp.BertHParams( epochs=3, batch_size=48, learning_rate=3e-5, - distribution_strategy='off')) + distribution_strategy='off', + ), + ) def test_predefined_average_word_embedding_spec(self): model_spec_obj = ( @@ -78,7 +80,7 @@ class ModelSpecTest(tf.test.TestCase): dropout_rate=0.2)) self.assertEqual( model_spec_obj.hparams, - hp.BaseHParams( + hp.AverageWordEmbeddingHParams( epochs=10, batch_size=32, learning_rate=0, @@ -101,7 +103,7 @@ class ModelSpecTest(tf.test.TestCase): custom_bert_classifier_options) def test_custom_average_word_embedding_spec(self): - custom_hparams = hp.BaseHParams( + custom_hparams = hp.AverageWordEmbeddingHParams( learning_rate=0.4, batch_size=64, epochs=10, @@ -110,7 +112,8 @@ class ModelSpecTest(tf.test.TestCase): export_dir='foo/bar', distribution_strategy='mirrored', num_gpus=3, - tpu='tpu/address') + tpu='tpu/address', + ) custom_average_word_embedding_model_options = ( classifier_model_options.AverageWordEmbeddingModelOptions( seq_len=512, diff --git a/mediapipe/model_maker/python/text/text_classifier/text_classifier.py b/mediapipe/model_maker/python/text/text_classifier/text_classifier.py index 59369931..6c8adc82 100644 --- a/mediapipe/model_maker/python/text/text_classifier/text_classifier.py +++ b/mediapipe/model_maker/python/text/text_classifier/text_classifier.py @@ -19,15 +19,16 @@ import tempfile from typing import Any, Optional, Sequence, Tuple import tensorflow as tf +from tensorflow_addons import optimizers as tfa_optimizers import tensorflow_hub as hub -from mediapipe.model_maker.python.core import hyperparameters as hp from mediapipe.model_maker.python.core.data import dataset as ds from mediapipe.model_maker.python.core.tasks import classifier from mediapipe.model_maker.python.core.utils import metrics from mediapipe.model_maker.python.core.utils import model_util from mediapipe.model_maker.python.core.utils import quantization from mediapipe.model_maker.python.text.text_classifier import dataset as text_ds +from mediapipe.model_maker.python.text.text_classifier import hyperparameters as hp from mediapipe.model_maker.python.text.text_classifier import model_options as mo from mediapipe.model_maker.python.text.text_classifier import model_spec as ms from mediapipe.model_maker.python.text.text_classifier import preprocessor @@ -55,22 +56,26 @@ def _validate(options: text_classifier_options.TextClassifierOptions): ms.SupportedModels.AVERAGE_WORD_EMBEDDING_CLASSIFIER)): raise ValueError("Expected AVERAGE_WORD_EMBEDDING_CLASSIFIER," f" got {options.supported_model}") - if (isinstance(options.model_options, mo.BertModelOptions) and - (options.supported_model != ms.SupportedModels.MOBILEBERT_CLASSIFIER)): + if isinstance(options.model_options, mo.BertModelOptions) and ( + options.supported_model != ms.SupportedModels.MOBILEBERT_CLASSIFIER + and options.supported_model != ms.SupportedModels.EXBERT_CLASSIFIER + ): raise ValueError( - f"Expected MOBILEBERT_CLASSIFIER, got {options.supported_model}") + "Expected a Bert Classifier(MobileBERT or EXBERT), got " + f"{options.supported_model}" + ) class TextClassifier(classifier.Classifier): """API for creating and training a text classification model.""" - def __init__(self, model_spec: Any, hparams: hp.BaseHParams, - label_names: Sequence[str]): + def __init__( + self, model_spec: Any, label_names: Sequence[str], shuffle: bool + ): super().__init__( - model_spec=model_spec, label_names=label_names, shuffle=hparams.shuffle) + model_spec=model_spec, label_names=label_names, shuffle=shuffle + ) self._model_spec = model_spec - self._hparams = hparams - self._callbacks = model_util.get_default_callbacks(self._hparams.export_dir) self._text_preprocessor: preprocessor.TextClassifierPreprocessor = None @classmethod @@ -107,7 +112,10 @@ class TextClassifier(classifier.Classifier): if options.hparams is None: options.hparams = options.supported_model.value().hparams - if options.supported_model == ms.SupportedModels.MOBILEBERT_CLASSIFIER: + if ( + options.supported_model == ms.SupportedModels.MOBILEBERT_CLASSIFIER + or options.supported_model == ms.SupportedModels.EXBERT_CLASSIFIER + ): text_classifier = ( _BertClassifier.create_bert_classifier(train_data, validation_data, options, @@ -225,11 +233,17 @@ class _AverageWordEmbeddingClassifier(TextClassifier): _DELIM_REGEX_PATTERN = r"[^\w\']+" - def __init__(self, model_spec: ms.AverageWordEmbeddingClassifierSpec, - model_options: mo.AverageWordEmbeddingModelOptions, - hparams: hp.BaseHParams, label_names: Sequence[str]): - super().__init__(model_spec, hparams, label_names) + def __init__( + self, + model_spec: ms.AverageWordEmbeddingClassifierSpec, + model_options: mo.AverageWordEmbeddingModelOptions, + hparams: hp.AverageWordEmbeddingHParams, + label_names: Sequence[str], + ): + super().__init__(model_spec, label_names, hparams.shuffle) self._model_options = model_options + self._hparams = hparams + self._callbacks = model_util.get_default_callbacks(self._hparams.export_dir) self._loss_function = "sparse_categorical_crossentropy" self._metric_functions = [ "accuracy", @@ -344,10 +358,16 @@ class _BertClassifier(TextClassifier): _INITIALIZER_RANGE = 0.02 - def __init__(self, model_spec: ms.BertClassifierSpec, - model_options: mo.BertModelOptions, hparams: hp.BaseHParams, - label_names: Sequence[str]): - super().__init__(model_spec, hparams, label_names) + def __init__( + self, + model_spec: ms.BertClassifierSpec, + model_options: mo.BertModelOptions, + hparams: hp.BertHParams, + label_names: Sequence[str], + ): + super().__init__(model_spec, label_names, hparams.shuffle) + self._hparams = hparams + self._callbacks = model_util.get_default_callbacks(self._hparams.export_dir) self._model_options = model_options with self._hparams.get_strategy().scope(): self._loss_function = tf.keras.losses.SparseCategoricalCrossentropy() @@ -480,11 +500,26 @@ class _BertClassifier(TextClassifier): initial_learning_rate=initial_lr, decay_schedule_fn=lr_schedule, warmup_steps=warmup_steps) - - self._optimizer = tf.keras.optimizers.experimental.AdamW( - lr_schedule, weight_decay=0.01, epsilon=1e-6, global_clipnorm=1.0) - self._optimizer.exclude_from_weight_decay( - var_names=["LayerNorm", "layer_norm", "bias"]) + if self._hparams.optimizer == hp.BertOptimizer.ADAMW: + self._optimizer = tf.keras.optimizers.experimental.AdamW( + lr_schedule, weight_decay=0.01, epsilon=1e-6, global_clipnorm=1.0 + ) + self._optimizer.exclude_from_weight_decay( + var_names=["LayerNorm", "layer_norm", "bias"] + ) + elif self._hparams.optimizer == hp.BertOptimizer.LAMB: + self._optimizer = tfa_optimizers.LAMB( + lr_schedule, + weight_decay_rate=0.01, + epsilon=1e-6, + exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"], + global_clipnorm=1.0, + ) + else: + raise ValueError( + "BertHParams.optimizer must be set to ADAM or " + f"LAMB. Got {self._hparams.optimizer}." + ) def _save_vocab(self, vocab_filepath: str): tf.io.gfile.copy( diff --git a/mediapipe/model_maker/python/text/text_classifier/text_classifier_demo.py b/mediapipe/model_maker/python/text/text_classifier/text_classifier_demo.py index c3d1711d..934bb1c4 100644 --- a/mediapipe/model_maker/python/text/text_classifier/text_classifier_demo.py +++ b/mediapipe/model_maker/python/text/text_classifier/text_classifier_demo.py @@ -66,14 +66,16 @@ def run(data_dir, quantization_config = None if (supported_model == text_classifier.SupportedModels.AVERAGE_WORD_EMBEDDING_CLASSIFIER): - hparams = text_classifier.HParams( - epochs=10, batch_size=32, learning_rate=0, export_dir=export_dir) + hparams = text_classifier.AverageWordEmbeddingHParams( + epochs=10, batch_size=32, learning_rate=0, export_dir=export_dir + ) # Warning: This takes extremely long to run on CPU elif ( supported_model == text_classifier.SupportedModels.MOBILEBERT_CLASSIFIER): quantization_config = quantization.QuantizationConfig.for_dynamic() - hparams = text_classifier.HParams( - epochs=3, batch_size=48, learning_rate=3e-5, export_dir=export_dir) + hparams = text_classifier.BertHParams( + epochs=3, batch_size=48, learning_rate=3e-5, export_dir=export_dir + ) # Fine-tunes the model. options = text_classifier.TextClassifierOptions( diff --git a/mediapipe/model_maker/python/text/text_classifier/text_classifier_options.py b/mediapipe/model_maker/python/text/text_classifier/text_classifier_options.py index c62fb27b..b61731f1 100644 --- a/mediapipe/model_maker/python/text/text_classifier/text_classifier_options.py +++ b/mediapipe/model_maker/python/text/text_classifier/text_classifier_options.py @@ -16,7 +16,7 @@ import dataclasses from typing import Optional -from mediapipe.model_maker.python.core import hyperparameters as hp +from mediapipe.model_maker.python.text.text_classifier import hyperparameters as hp from mediapipe.model_maker.python.text.text_classifier import model_options as mo from mediapipe.model_maker.python.text.text_classifier import model_spec as ms @@ -34,5 +34,5 @@ class TextClassifierOptions: architecture of the `supported_model`. """ supported_model: ms.SupportedModels - hparams: Optional[hp.BaseHParams] = None + hparams: Optional[hp.HParams] = None model_options: Optional[mo.TextClassifierModelOptions] = None diff --git a/mediapipe/model_maker/python/text/text_classifier/text_classifier_test.py b/mediapipe/model_maker/python/text/text_classifier/text_classifier_test.py index 34830c9f..e6057059 100644 --- a/mediapipe/model_maker/python/text/text_classifier/text_classifier_test.py +++ b/mediapipe/model_maker/python/text/text_classifier/text_classifier_test.py @@ -66,12 +66,14 @@ class TextClassifierTest(tf.test.TestCase): def test_create_and_train_average_word_embedding_model(self): train_data, validation_data = self._get_data() - options = ( - text_classifier.TextClassifierOptions( - supported_model=(text_classifier.SupportedModels - .AVERAGE_WORD_EMBEDDING_CLASSIFIER), - hparams=text_classifier.HParams( - epochs=1, batch_size=1, learning_rate=0))) + options = text_classifier.TextClassifierOptions( + supported_model=( + text_classifier.SupportedModels.AVERAGE_WORD_EMBEDDING_CLASSIFIER + ), + hparams=text_classifier.AverageWordEmbeddingHParams( + epochs=1, batch_size=1, learning_rate=0 + ), + ) average_word_embedding_classifier = ( text_classifier.TextClassifier.create(train_data, validation_data, options)) @@ -103,12 +105,15 @@ class TextClassifierTest(tf.test.TestCase): options = text_classifier.TextClassifierOptions( supported_model=text_classifier.SupportedModels.MOBILEBERT_CLASSIFIER, model_options=text_classifier.BertModelOptions( - do_fine_tuning=False, seq_len=2), - hparams=text_classifier.HParams( + do_fine_tuning=False, seq_len=2 + ), + hparams=text_classifier.BertHParams( epochs=1, batch_size=1, learning_rate=3e-5, - distribution_strategy='off')) + distribution_strategy='off', + ), + ) bert_classifier = text_classifier.TextClassifier.create( train_data, validation_data, options)