def create(train_data,
model_spec,
batch_size=None,
epochs=2,
shuffle=False,
do_train=True):
"""Loads data and train the model for question answer.
Args:
train_data: Training data.
model_spec: Specification for the model.
batch_size: Batch size for training.
epochs: Number of epochs for training.
shuffle: Whether the data should be shuffled.
do_train: Whether to run training.
Returns:
object of QuestionAnswer class.
"""
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
model = QuestionAnswer(model_spec, shuffle=shuffle)
if do_train:
tf.compat.v1.logging.info('Retraining the models...')
model.train(train_data, epochs, batch_size)
else:
model.create_model()
return model
def create(train_data,
model_spec,
validation_data=None,
epochs=None,
batch_size=None,
do_train=True):
"""Loads data and train the model for object detection.
Args:
train_data: Training data.
model_spec: Specification for the model.
validation_data: Validation data. If None, skips validation process.
epochs: Number of epochs for training.
batch_size: Batch size for training.
do_train: Whether to run training.
Returns:
ObjectDetector
"""
model_spec = ms.get(model_spec)
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
object_detector = ObjectDetector(model_spec, train_data.label_map)
if do_train:
tf.compat.v1.logging.info('Retraining the models...')
object_detector.train(train_data, validation_data, epochs, batch_size)
else:
object_detector.create_model()
return object_detector
def create(train_data,
model_spec=ms.AverageWordVecModelSpec(),
shuffle=False,
batch_size=32,
epochs=None,
validation_data=None):
"""Loads data and train the model for test classification.
Args:
train_data: Training data.
model_spec: Specification for the model.
shuffle: Whether the data should be shuffled.
batch_size: Batch size for training.
epochs: Number of epochs for training.
validation_data: Validation data. If None, skips validation process.
Returns:
TextClassifier
"""
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError('Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
text_classifier = TextClassifier(
model_spec,
train_data.index_to_label,
train_data.num_classes,
shuffle=shuffle)
tf.compat.v1.logging.info('Retraining the models...')
text_classifier.train(train_data, validation_data, epochs, batch_size)
return text_classifier
def export_tflite(model,
tflite_filepath,
quantization_config=None,
convert_from_saved_model_tf2=False,
preprocess=None,
supported_ops=(tf.lite.OpsSet.TFLITE_BUILTINS, )):
"""Converts the retrained model to tflite format and saves it.
Args:
model: model to be converted to tflite.
tflite_filepath: File path to save tflite model.
quantization_config: Configuration for post-training quantization.
convert_from_saved_model_tf2: Convert to TFLite from saved_model in TF 2.x.
preprocess: A preprocess function to apply on the dataset.
# TODO(wangtz): Remove when preprocess is split off from CustomModel.
supported_ops: A list of supported ops in the converted TFLite file.
"""
if tflite_filepath is None:
raise ValueError(
"TFLite filepath couldn't be None when exporting to tflite.")
if compat.get_tf_behavior() == 1:
lite = tf.compat.v1.lite
else:
lite = tf.lite
convert_from_saved_model = (compat.get_tf_behavior() == 1
or convert_from_saved_model_tf2)
with _create_temp_dir(convert_from_saved_model) as temp_dir_name:
if temp_dir_name:
save_path = os.path.join(temp_dir_name, 'saved_model')
model.save(save_path, include_optimizer=False, save_format='tf')
converter = lite.TFLiteConverter.from_saved_model(save_path)
else:
converter = lite.TFLiteConverter.from_keras_model(model)
# TODO(b/191205988): Explicitly disable saved model lowering in
# the conversion.
converter.experimental_lower_to_saved_model = False
if quantization_config:
converter = quantization_config.get_converter_with_quantization(
converter, preprocess=preprocess)
converter.target_spec.supported_ops = supported_ops
tflite_model = converter.convert()
with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
f.write(tflite_model)
def _get_model_info(model_spec, num_classes, quantized=False, version='v1'):
"""Gets the specific info for the image model."""
if not isinstance(model_spec, ms.ImageModelSpec):
raise ValueError(
'Currently only support models for image classification.')
name = model_spec.name
if quantized:
name += '_quantized'
if quantized and compat.get_tf_behavior() == 1:
image_min = 0
image_max = 255
else:
image_min = 0
image_max = 1
return metadata_writer.ModelSpecificInfo(
model_spec.name,
version,
image_width=model_spec.input_image_shape[1],
image_height=model_spec.input_image_shape[0],
mean=model_spec.mean_rgb,
std=model_spec.stddev_rgb,
image_min=image_min,
image_max=image_max,
num_classes=num_classes)
def _test_export_to_tflite(self, model):
tflite_output_file = os.path.join(self.get_temp_dir(), 'model.tflite')
labels_output_file = os.path.join(self.get_temp_dir(), 'label')
model.export(tflite_output_file, labels_output_file)
labels = self._load_labels(labels_output_file)
self.assertEqual(labels, ['cyan', 'magenta', 'yellow'])
lite_model = self._load_lite_model(tflite_output_file)
test_ds = model._gen_dataset(
self.test_data, batch_size=1, is_training=False)
if compat.get_tf_behavior() == 1:
iterator = test_ds.make_one_shot_iterator()
image_tensor, label_tensor = iterator.get_next()
with tf.compat.v1.Session() as sess:
for _ in range(self.test_data.size):
image, label = sess.run((image_tensor, label_tensor))
output_batch = lite_model(image)
prediction = np.argmax(output_batch[0])
label = np.argmax(label[0])
self.assertEqual(label, prediction)
else:
for image, label in test_ds:
output_batch = lite_model(image.numpy())
prediction = np.argmax(output_batch[0])
label = np.argmax(label.numpy()[0])
self.assertEqual(label, prediction)
def _test_export_to_tflite(self, model):
tflite_output_file = os.path.join(self.get_temp_dir(), 'model.tflite')
labels_output_file = os.path.join(self.get_temp_dir(), 'label')
model.export(tflite_output_file, labels_output_file)
labels = self._load_labels(labels_output_file)
self.assertEqual(labels, ['cyan', 'magenta', 'yellow'])
lite_model = self._load_lite_model(tflite_output_file)
if compat.get_tf_behavior() == 1:
image_placeholder = tf.compat.v1.placeholder(
tf.uint8, [1, self.IMAGE_SIZE, self.IMAGE_SIZE, 3])
label_placeholder = tf.compat.v1.placeholder(tf.int32, [1])
image_tensor, _ = model.preprocess(image_placeholder, label_placeholder)
with tf.compat.v1.Session() as sess:
for i, (class_name, rgb) in enumerate(self.CMY_NAMES_AND_RGB_VALUES):
input_image = np.expand_dims(_fill_image(rgb, self.IMAGE_SIZE), 0)
image = sess.run(
image_tensor,
feed_dict={
image_placeholder: input_image,
label_placeholder: [i]
})
output_batch = lite_model(image)
prediction = labels[np.argmax(output_batch[0])]
self.assertEqual(class_name, prediction)
else:
for i, (class_name, rgb) in enumerate(self.CMY_NAMES_AND_RGB_VALUES):
input_batch = np.expand_dims(_fill_image(rgb, self.IMAGE_SIZE), 0)
image, _ = model.preprocess(input_batch, i)
image = image.numpy()
output_batch = lite_model(image)
prediction = labels[np.argmax(output_batch[0])]
self.assertEqual(class_name, prediction)
def create(cls,
train_data: object_detector_dataloader.DataLoader,
model_spec: object_detector_spec.EfficientDetModelSpec,
validation_data: Optional[
object_detector_dataloader.DataLoader] = None,
epochs: Optional[object_detector_dataloader.DataLoader] = None,
batch_size: Optional[int] = None,
train_whole_model: bool = False,
do_train: bool = True) -> T:
"""Loads data and train the model for object detection.
Args:
train_data: Training data.
model_spec: Specification for the model.
validation_data: Validation data. If None, skips validation process.
epochs: Number of epochs for training.
batch_size: Batch size for training.
train_whole_model: Boolean, False by default. If true, train the whole
model. Otherwise, only train the layers that are not match
`model_spec.config.var_freeze_expr`.
do_train: Whether to run training.
Returns:
An instance based on ObjectDetector.
"""
model_spec = ms.get(model_spec)
if epochs is not None:
model_spec.config.num_epochs = epochs
if batch_size is not None:
model_spec.config.batch_size = batch_size
if train_whole_model:
model_spec.config.var_freeze_expr = None
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
object_detector = cls(model_spec, train_data.label_map, train_data)
if do_train:
tf.compat.v1.logging.info('Retraining the models...')
object_detector.train(train_data, validation_data, epochs,
batch_size)
else:
object_detector.create_model()
return object_detector
def export_tflite(model,
tflite_filepath,
quantization_config=None,
gen_dataset_fn=None,
convert_from_saved_model_tf2=False):
"""Converts the retrained model to tflite format and saves it.
Args:
model: model to be converted to tflite.
tflite_filepath: File path to save tflite model.
quantization_config: Configuration for post-training quantization.
gen_dataset_fn: Function to generate tf.data.dataset from
`representative_data`. Used only when `representative_data` in
`quantization_config` is setted.
convert_from_saved_model_tf2: Convert to TFLite from saved_model in TF 2.x.
"""
if tflite_filepath is None:
raise ValueError(
"TFLite filepath couldn't be None when exporting to tflite.")
if compat.get_tf_behavior() == 1:
lite = tf.compat.v1.lite
else:
lite = tf.lite
convert_from_saved_model = (compat.get_tf_behavior() == 1
or convert_from_saved_model_tf2)
with _create_temp_dir(True) as temp_dir_name:
if convert_from_saved_model:
save_path = os.path.join(temp_dir_name, 'saved_model')
model.save(save_path, include_optimizer=False, save_format='tf')
converter = lite.TFLiteConverter.from_saved_model(save_path)
else:
# TODO modified by [email protected]
save_path = os.path.join(temp_dir_name, 'keras_model.h5')
tf.keras.experimental.export_saved_model(model, save_path)
converter = lite.TFLiteConverter.from_saved_model(save_path)
if quantization_config:
converter = quantization_config.get_converter_with_quantization(
converter, gen_dataset_fn)
tflite_model = converter.convert()
with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
f.write(tflite_model)
def _export_tflite(self,
tflite_filename,
label_filename,
quantized=False,
quantization_steps=None,
representative_data=None):
"""Converts the retrained model to tflite format and saves it.
Args:
tflite_filename: File name to save tflite model.
label_filename: File name to save labels.
quantized: boolean, if True, save quantized model.
quantization_steps: Number of post-training quantization calibration steps
to run. Used only if `quantized` is True.
representative_data: Representative data used for post-training
quantization. Used only if `quantized` is True.
"""
if compat.get_tf_behavior() == 1:
with tempfile.TemporaryDirectory() as temp_dir:
save_path = os.path.join(temp_dir, 'saved_model')
self.model.save(save_path,
include_optimizer=False,
save_format='tf')
converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(
save_path)
else:
converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
if quantized:
if quantization_steps is None:
quantization_steps = DEFAULT_QUANTIZATION_STEPS
if representative_data is None:
raise ValueError(
'representative_data couldn\'t be None if model is quantized.'
)
ds = self._gen_dataset(representative_data,
batch_size=1,
is_training=False)
converter.representative_dataset = tf.lite.RepresentativeDataset(
get_representative_dataset_gen(ds, quantization_steps))
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
tflite_model = converter.convert()
with tf.io.gfile.GFile(tflite_filename, 'wb') as f:
f.write(tflite_model)
with tf.io.gfile.GFile(label_filename, 'w') as f:
f.write('\n'.join(self.index_to_label))
tf.compat.v1.logging.info(
'Export to tflite model %s, saved labels in %s.', tflite_filename,
label_filename)
def create(cls,
train_data,
model_spec='average_word_vec',
validation_data=None,
batch_size=None,
epochs=3,
steps_per_epoch=None,
shuffle=False,
do_train=True):
"""Loads data and train the model for test classification.
Args:
train_data: Training data.
model_spec: Specification for the model.
validation_data: Validation data. If None, skips validation process.
batch_size: Batch size for training.
epochs: Number of epochs for training.
steps_per_epoch: Integer or None. Total number of steps (batches of
samples) before declaring one epoch finished and starting the next
epoch. If `steps_per_epoch` is None, the epoch will run until the input
dataset is exhausted.
shuffle: Whether the data should be shuffled.
do_train: Whether to run training.
Returns:
An instance based on TextClassifier.
"""
model_spec = ms.get(model_spec)
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
text_classifier = cls(model_spec,
train_data.index_to_label,
shuffle=shuffle)
if do_train:
tf.compat.v1.logging.info('Retraining the models...')
text_classifier.train(train_data, validation_data, epochs,
batch_size, steps_per_epoch)
else:
text_classifier.create_model()
return text_classifier
def create(cls,
train_data,
model_spec='average_word_vec',
validation_data=None,
batch_size=None,
epochs=3,
shuffle=False,
do_train=True):
"""Loads data and train the model for test classification.
Args:
train_data: Training data.
model_spec: Specification for the model.
validation_data: Validation data. If None, skips validation process.
batch_size: Batch size for training.
epochs: Number of epochs for training.
shuffle: Whether the data should be shuffled.
do_train: Whether to run training.
Returns:
An instance based on TextClassifier.
"""
model_spec = ms.get(model_spec)
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
text_classifier = cls(model_spec,
train_data.index_to_label,
shuffle=shuffle)
if do_train:
tf.compat.v1.logging.info('Retraining the models...')
text_classifier.train(train_data, validation_data, epochs,
batch_size)
else:
text_classifier.create_model()
return text_classifier
def generate_elements(ds):
"""Generates elements from `tf.data.dataset`."""
if compat.get_tf_behavior() == 2:
for element in ds.as_numpy_iterator():
yield element
else:
iterator = ds.make_one_shot_iterator()
next_element = iterator.get_next()
with tf.compat.v1.Session() as sess:
try:
while True:
yield sess.run(next_element)
except tf.errors.OutOfRangeError:
return
def representative_dataset_gen():
"""Generates representative dataset for quantized."""
if compat.get_tf_behavior() == 2:
for image, _ in dataset.take(num_steps):
yield [image]
else:
iterator = tf.compat.v1.data.make_one_shot_iterator(
dataset.take(num_steps))
next_element = iterator.get_next()
with tf.compat.v1.Session() as sess:
while True:
try:
image, _ = sess.run(next_element)
yield [image]
except tf.errors.OutOfRangeError:
break
def _get_model_info(model_spec,
num_classes,
quantization_config=None,
version='v1'):
"""Gets the specific info for the image model."""
if not isinstance(model_spec, ms.ImageModelSpec):
raise ValueError(
'Currently only support models for image classification.')
image_min = 0
image_max = 1
name = model_spec.name
if quantization_config:
name += '_quantized'
# TODO(yuqili): Remove `compat.get_tf_behavior() == 1` once b/153576655 is
# fixed.
if compat.get_tf_behavior() == 1:
if quantization_config.inference_input_type == tf.uint8:
image_min = 0
image_max = 255
elif quantization_config.inference_input_type == tf.int8:
image_min = -128
image_max = 127
return metadata_writer.ModelSpecificInfo(
model_spec.name,
version,
image_width=model_spec.input_image_shape[1],
image_height=model_spec.input_image_shape[0],
mean=model_spec.mean_rgb,
std=model_spec.stddev_rgb,
image_min=image_min,
image_max=image_max,
num_classes=num_classes)
def create_full_integer_quantization(
cls,
representative_data,
quantization_steps=DEFAULT_QUANTIZATION_STEPS,
optimizations=tf.lite.Optimize.DEFAULT,
inference_input_type=tf.uint8,
inference_output_type=tf.uint8,
is_integer_only=False):
"""Creates configuration for full integer quantization.
Args:
representative_data: Representative data used for post-training
quantization.
quantization_steps: Number of post-training quantization calibration steps
to run.
optimizations: A list of optimizations to apply when converting the model.
If not set, use `[Optimize.DEFAULT]` by default.
inference_input_type: Target data type of real-number input arrays. Used
only when `is_integer_only` is True. For TensorFlow 2, it is set to
`tf.float32`. For TensorFlow 1, it must be in `{tf.uint8, tf.int8}`.
inference_output_type: Target data type of real-number output arrays. Used
only when `is_integer_only` is True. For TensorFlow 2, it is set to
`tf.float32`. For TensorFlow 1, it must be in `{tf.uint8, tf.int8}`.
is_integer_only: If True, enforces full integer quantization for all ops
including the input and output. If False, uses integer with float
fallback (using default float input/output) that mean to fully integer
quantize a model, but use float operators when they don't have an
integer implementation.
Returns:
QuantizationConfig.
"""
if not is_integer_only:
return QuantizationConfig(optimizations,
representative_data=representative_data,
quantization_steps=quantization_steps)
else:
if compat.get_tf_behavior() == 2:
# TODO(b/153576655): Replicate inference_input_type and
# inference_output_type flags in TFLiteConverterV2
tf.compat.v1.logging.warning(
'For integer only quantization, `inference_input_type` and '
'`inference_output_type` are set to tf.float32. Support for '
'tf.int8 and tf.uint8 will be added soon.')
inference_input_type = tf.float32
inference_output_type = tf.float32
else:
if inference_input_type not in [tf.uint8, tf.int8]:
raise ValueError('For integer only quantization, '
'`inference_input_type` '
'should be tf.int8 or tf.uint8.')
if inference_output_type not in [tf.uint8, tf.int8]:
raise ValueError('For integer only quantization, '
'`inference_output_type` '
'should be tf.int8 or tf.uint8.')
return QuantizationConfig(
optimizations,
representative_data=representative_data,
quantization_steps=quantization_steps,
inference_input_type=inference_input_type,
inference_output_type=inference_output_type,
supported_ops=[tf.lite.OpsSet.TFLITE_BUILTINS_INT8])
def create(train_data,
model_export_format=mef.ModelExportFormat.TFLITE,
model_spec=ms.efficientnet_lite0_spec,
shuffle=False,
validation_data=None,
batch_size=None,
epochs=None,
train_whole_model=None,
dropout_rate=None,
learning_rate=None,
momentum=None):
"""Loads data and retrains the model based on data for image classification.
Args:
train_data: Training data.
model_export_format: Model export format such as saved_model / tflite.
model_spec: Specification for the model.
shuffle: Whether the data should be shuffled.
validation_data: Validation data. If None, skips validation process.
batch_size: Number of samples per training step.
epochs: Number of epochs for training.
train_whole_model: If true, the Hub module is trained together with the
classification layer on top. Otherwise, only train the top classification
layer.
dropout_rate: the rate for dropout.
learning_rate: a Python float forwarded to the optimizer.
momentum: a Python float forwarded to the optimizer.
Returns:
An instance of ImageClassifier class.
"""
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
# The hyperparameters for make_image_classifier by tensorflow hub.
hparams = lib.get_default_hparams()
if batch_size is not None:
hparams = hparams._replace(batch_size=batch_size)
if epochs is not None:
hparams = hparams._replace(train_epochs=epochs)
if train_whole_model is not None:
hparams = hparams._replace(do_fine_tuning=train_whole_model)
if dropout_rate is not None:
hparams = hparams._replace(dropout_rate=dropout_rate)
if learning_rate is not None:
hparams = hparams._replace(learning_rate=learning_rate)
if momentum is not None:
hparams = hparams._replace(momentum=momentum)
image_classifier = ImageClassifier(model_export_format,
model_spec,
train_data.index_to_label,
train_data.num_classes,
shuffle=shuffle,
hparams=hparams)
tf.compat.v1.logging.info('Retraining the models...')
image_classifier.train(train_data, validation_data)
return image_classifier
def decorator(*args, **kwargs):
if compat.get_tf_behavior() not in [1, 2]:
tf.compat.v1.logging.info(
"Skip function {} for test_in_tf_1and2".format(fn.__name__))
return
fn(*args, **kwargs)
def _export_tflite(self,
tflite_filepath,
quantized=False,
quantization_steps=None,
representative_data=None,
inference_input_type=tf.float32,
inference_output_type=tf.float32):
"""Converts the retrained model to tflite format and saves it.
Args:
tflite_filepath: File path to save tflite model.
quantized: boolean, if True, save quantized model.
quantization_steps: Number of post-training quantization calibration steps
to run. Used only if `quantized` is True.
representative_data: Representative data used for post-training
quantization. Used only if `quantized` is True.
inference_input_type: Target data type of real-number input arrays. Allows
for a different type for input arrays. Defaults to tf.float32. Must be
be `{tf.float32, tf.uint8, tf.int8}`
inference_output_type: Target data type of real-number output arrays.
Allows for a different type for output arrays. Defaults to tf.float32.
Must be `{tf.float32, tf.uint8, tf.int8}`
"""
if tflite_filepath is None:
raise ValueError(
"TFLite filepath couldn't be None when exporting to tflite.")
tf.compat.v1.logging.info('Exporting to tflite model in %s.',
tflite_filepath)
temp_dir = None
if compat.get_tf_behavior() == 1:
temp_dir = tempfile.TemporaryDirectory()
save_path = os.path.join(temp_dir.name, 'saved_model')
self.model.save(save_path,
include_optimizer=False,
save_format='tf')
converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(
save_path)
else:
converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
if quantized:
if quantization_steps is None:
quantization_steps = DEFAULT_QUANTIZATION_STEPS
if representative_data is None:
raise ValueError(
'representative_data couldn\'t be None if model is quantized.'
)
ds = self._gen_dataset(representative_data,
batch_size=1,
is_training=False)
converter.representative_dataset = tf.lite.RepresentativeDataset(
get_representative_dataset_gen(ds, quantization_steps))
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.inference_input_type = inference_input_type
converter.inference_output_type = inference_output_type
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
tflite_model = converter.convert()
if temp_dir:
temp_dir.cleanup()
with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
f.write(tflite_model)
def __init__(
self,
uri='https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1',
model_dir=None,
seq_len=128,
dropout_rate=0.1,
initializer_range=0.02,
learning_rate=3e-5,
distribution_strategy='mirrored',
num_gpus=-1,
tpu='',
trainable=True,
do_lower_case=True,
is_tf2=True,
convert_from_saved_model_tf2=False):
"""Initialze an instance with model paramaters.
Args:
uri: TF-Hub path/url to Bert module.
model_dir: The location of the model checkpoint files.
seq_len: Length of the sequence to feed into the model.
dropout_rate: The rate for dropout.
initializer_range: The stdev of the truncated_normal_initializer for
initializing all weight matrices.
learning_rate: The initial learning rate for Adam.
distribution_strategy: A string specifying which distribution strategy to
use. Accepted values are 'off', 'one_device', 'mirrored',
'parameter_server', 'multi_worker_mirrored', and 'tpu' -- case
insensitive. 'off' means not to use Distribution Strategy; 'tpu' means
to use TPUStrategy using `tpu_address`.
num_gpus: How many GPUs to use at each worker with the
DistributionStrategies API. The default is -1, which means utilize all
available GPUs.
tpu: TPU address to connect to.
trainable: boolean, whether pretrain layer is trainable.
do_lower_case: boolean, whether to lower case the input text. Should be
True for uncased models and False for cased models.
is_tf2: boolean, whether the hub module is in TensorFlow 2.x format.
convert_from_saved_model_tf2: Convert to TFLite from saved_model in TF
2.x.
"""
if compat.get_tf_behavior() not in self.compat_tf_versions:
raise ValueError('Incompatible versions. Expect {}, but got {}.'.format(
self.compat_tf_versions, compat.get_tf_behavior()))
self.seq_len = seq_len
self.dropout_rate = dropout_rate
self.initializer_range = initializer_range
self.learning_rate = learning_rate
self.trainable = trainable
self.model_dir = model_dir
if self.model_dir is None:
self.model_dir = tempfile.mkdtemp()
num_gpus = get_num_gpus(num_gpus)
self.strategy = distribution_utils.get_distribution_strategy(
distribution_strategy=distribution_strategy,
num_gpus=num_gpus,
tpu_address=tpu)
self.tpu = tpu
self.uri = uri
self.do_lower_case = do_lower_case
self.is_tf2 = is_tf2
self.bert_config = bert_configs.BertConfig(
0,
initializer_range=self.initializer_range,
hidden_dropout_prob=self.dropout_rate)
self.convert_from_saved_model_tf2 = convert_from_saved_model_tf2
self.is_built = False
def create(train_data,
model_spec='efficientnet_lite0',
validation_data=None,
batch_size=None,
epochs=None,
train_whole_model=None,
dropout_rate=None,
learning_rate=None,
momentum=None,
shuffle=False,
use_augmentation=False,
use_hub_library=True,
warmup_steps=None,
model_dir=None,
do_train=True):
"""Loads data and retrains the model based on data for image classification.
Args:
train_data: Training data.
model_spec: Specification for the model.
validation_data: Validation data. If None, skips validation process.
batch_size: Number of samples per training step. If `use_hub_library` is
False, it represents the base learning rate when train batch size is 256
and it's linear to the batch size.
epochs: Number of epochs for training.
train_whole_model: If true, the Hub module is trained together with the
classification layer on top. Otherwise, only train the top classification
layer.
dropout_rate: The rate for dropout.
learning_rate: Base learning rate when train batch size is 256. Linear to
the batch size.
momentum: a Python float forwarded to the optimizer. Only used when
`use_hub_library` is True.
shuffle: Whether the data should be shuffled.
use_augmentation: Use data augmentation for preprocessing.
use_hub_library: Use `make_image_classifier_lib` from tensorflow hub to
retrain the model.
warmup_steps: Number of warmup steps for warmup schedule on learning rate.
If None, the default warmup_steps is used which is the total training
steps in two epochs. Only used when `use_hub_library` is False.
model_dir: The location of the model checkpoint files. Only used when
`use_hub_library` is False.
do_train: Whether to run training.
Returns:
An instance of ImageClassifier class.
"""
model_spec = ms.get(model_spec)
if compat.get_tf_behavior() not in model_spec.compat_tf_versions:
raise ValueError(
'Incompatible versions. Expect {}, but got {}.'.format(
model_spec.compat_tf_versions, compat.get_tf_behavior()))
if use_hub_library:
hparams = get_hub_lib_hparams(batch_size=batch_size,
train_epochs=epochs,
do_fine_tuning=train_whole_model,
dropout_rate=dropout_rate,
learning_rate=learning_rate,
momentum=momentum)
else:
hparams = train_image_classifier_lib.HParams.get_hparams(
batch_size=batch_size,
train_epochs=epochs,
do_fine_tuning=train_whole_model,
dropout_rate=dropout_rate,
learning_rate=learning_rate,
warmup_steps=warmup_steps,
model_dir=model_dir)
image_classifier = ImageClassifier(model_spec,
train_data.index_to_label,
shuffle=shuffle,
hparams=hparams,
use_augmentation=use_augmentation,
representative_data=train_data)
if do_train:
tf.compat.v1.logging.info('Retraining the models...')
image_classifier.train(train_data, validation_data)
else:
# Used in evaluation.
image_classifier.create_model(with_loss_and_metrics=True)
return image_classifier
def export_tflite(model,
tflite_filepath,
quantization_config=None,
convert_from_saved_model_tf2=False,
preprocess=None,
supported_ops=(tf.lite.OpsSet.TFLITE_BUILTINS, ),
quantization_type: Optional[QuantizationType] = None,
representative_dataset: Optional[tf.data.Dataset] = None):
"""Converts the retrained model to tflite format and saves it.
Args:
model: model to be converted to tflite.
tflite_filepath: File path to save tflite model.
quantization_config: Configuration for post-training quantization.
convert_from_saved_model_tf2: Convert to TFLite from saved_model in TF 2.x.
preprocess: A preprocess function to apply on the dataset.
# TODO(wangtz): Remove when preprocess is split off from CustomModel.
supported_ops: A list of supported ops in the converted TFLite file.
quantization_type: Enum, type of post-training quantization. Accepted values
are `INT8`, `FP16`, `FP32`, `DYNAMIC`. `FP16` means float16 quantization
with 2x smaller, optimized for GPU. `INT8` means full integer quantization
with 4x smaller, 3x+ speedup, optimized for EdgeTPU. 'DYNAMIC' means
dynamic range quantization with 4x smaller, 2x-3x speedup. `FP32` mean
exporting float model without quantization. Please refer to
https://www.tensorflow.org/lite/performance/post_training_quantization
for more detailed about different techniques for post-training
quantization.
representative_dataset: Representative dataset for full integer
quantization. Used when `quantization_type=INT8`.
"""
if quantization_type and quantization_config:
raise ValueError(
'At most one of the paramaters `quantization_type` and '
'`quantization_config` can be set.')
if tflite_filepath is None:
raise ValueError(
"TFLite filepath couldn't be None when exporting to tflite.")
if compat.get_tf_behavior() == 1:
lite = tf.compat.v1.lite
else:
lite = tf.lite
convert_from_saved_model = (compat.get_tf_behavior() == 1
or convert_from_saved_model_tf2)
with _create_temp_dir(convert_from_saved_model) as temp_dir_name:
if temp_dir_name:
save_path = os.path.join(temp_dir_name, 'saved_model')
model.save(save_path, include_optimizer=False, save_format='tf')
converter = lite.TFLiteConverter.from_saved_model(save_path)
else:
converter = lite.TFLiteConverter.from_keras_model(model)
# Uses `quantization_type` for post-training quantization.
if quantization_type == QuantizationType.INT8:
if representative_dataset is None:
raise ValueError('`representative_data` must be set when '
'`quantization_type=QuantizationType.INT8.')
converter.representative_dataset = configs.get_representative_dataset_gen(
representative_dataset, _NUM_CALIBRATION_STEPS)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.target_spec.supported_ops = supported_ops
elif quantization_type == QuantizationType.FP16:
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
elif quantization_type == QuantizationType.DYNAMIC:
converter.optimizations = [tf.lite.Optimize.DEFAULT]
elif quantization_type and quantization_type != QuantizationType.FP32:
raise ValueError('Unsupported `quantization_type`: %s' %
str(quantization_type))
# Uses `quantization_config` for post-training quantization.
elif quantization_config:
converter = quantization_config.get_converter_with_quantization(
converter, preprocess=preprocess)
converter.target_spec.supported_ops = supported_ops
tflite_model = converter.convert()
with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
f.write(tflite_model)