def __init__(self,
data,
model_export_format,
model_spec,
shuffle=True,
validation_ratio=0.1,
test_ratio=0.1,
hparams=lib.get_default_hparams()):
"""Init function for ImageClassifier class.
Including splitting the raw input data into train/eval/test sets and
selecting the exact NN model to be used.
Args:
data: Raw data that could be splitted for training / validation / testing.
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_ratio: The ratio of validation data to be splitted.
test_ratio: The ratio of test data to be splitted.
hparams: A namedtuple of hyperparameters. This function expects
.dropout_rate: The fraction of the input units to drop, used in dropout
layer.
"""
super(ImageClassifier,
self).__init__(data, model_export_format, model_spec, shuffle,
hparams.do_fine_tuning, validation_ratio, test_ratio)
self.hparams = hparams
self.model = self._create_model()
def train(self, hparams=lib.get_default_hparams()):
"""Feeds the training data for training.
Args:
hparams: A namedtuple of hyperparameters. This function expects
.train_epochs: a Python integer with the number of passes over the
training dataset;
.learning_rate: a Python float forwarded to the optimizer;
.momentum: a Python float forwarded to the optimizer;
.batch_size: a Python integer, number of samples per training step.
Returns:
The tf.keras.callbacks.History object returned by tf.keras.Model.fit*().
"""
train_data_and_size = (self._gen_train_dataset(self.train_data,
hparams.batch_size),
self.train_data.size)
validation_data_and_size = (self._gen_valid_dataset(self.valid_data,
hparams.batch_size),
self.valid_data.size)
# Trains the models.
return lib.train_model(self.model, hparams, train_data_and_size,
validation_data_and_size)
def __init__(self,
model_spec,
index_to_label,
shuffle=True,
hparams=hub_lib.get_default_hparams(),
use_augmentation=False,
representative_data=None):
"""Init function for ImageClassifier class.
Args:
model_spec: Specification for the model.
index_to_label: A list that map from index to label class name.
shuffle: Whether the data should be shuffled.
hparams: A namedtuple of hyperparameters. This function expects
.dropout_rate: The fraction of the input units to drop, used in dropout
layer.
.do_fine_tuning: If true, the Hub module is trained together with the
classification layer on top.
use_augmentation: Use data augmentation for preprocessing.
representative_data: Representative dataset for full integer
quantization. Used when converting the keras model to the TFLite model
with full interger quantization.
"""
super(ImageClassifier, self).__init__(model_spec, index_to_label,
shuffle, hparams.do_fine_tuning)
num_classes = len(index_to_label)
self._hparams = hparams
self.preprocess = image_preprocessing.Preprocessor(
self.model_spec.input_image_shape,
num_classes,
self.model_spec.mean_rgb,
self.model_spec.stddev_rgb,
use_augmentation=use_augmentation)
self.history = None # Training history that returns from `keras_model.fit`.
self.representative_data = representative_data
def __init__(self,
model_spec,
index_to_label,
num_classes,
shuffle=True,
hparams=hub_lib.get_default_hparams(),
use_augmentation=False):
"""Init function for ImageClassifier class.
Args:
model_spec: Specification for the model.
index_to_label: A list that map from index to label class name.
num_classes: Number of label classes.
shuffle: Whether the data should be shuffled.
hparams: A namedtuple of hyperparameters. This function expects
.dropout_rate: The fraction of the input units to drop, used in dropout
layer.
.do_fine_tuning: If true, the Hub module is trained together with the
classification layer on top.
use_augmentation: Use data augmentation for preprocessing.
"""
super(ImageClassifier,
self).__init__(model_spec, index_to_label, num_classes, shuffle,
hparams.do_fine_tuning)
self.hparams = hparams
self.model = self._create_model()
self.preprocessor = image_preprocessing.Preprocessor(
self.model_spec.input_image_shape,
num_classes,
self.model_spec.mean_rgb,
self.model_spec.stddev_rgb,
use_augmentation=use_augmentation)
self.history = None # Training history that returns from `keras_model.fit`.
def create(data,
model_export_format=mef.ModelExportFormat.TFLITE,
model_spec=ms.mobilenet_v2_spec,
shuffle=False,
validation_ratio=0.1,
test_ratio=0.1,
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:
data: Raw data that could be splitted for training / validation / testing.
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_ratio: The ratio of validation data to be splitted.
test_ratio: The ratio of test data to be splitted.
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.
"""
# 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(
data,
model_export_format,
model_spec,
shuffle=shuffle,
validation_ratio=validation_ratio,
test_ratio=test_ratio,
hparams=hparams)
tf.compat.v1.logging.info('Retraining the models...')
image_classifier.train()
return image_classifier
def __init__(self,
model_export_format,
model_spec,
index_to_label,
num_classes,
shuffle=True,
hparams=lib.get_default_hparams(),
use_augmentation=False):
"""Init function for ImageClassifier class.
Args:
model_export_format: Model export format such as saved_model / tflite.
model_spec: Specification for the model.
index_to_label: A list that map from index to label class name.
num_classes: Number of label classes.
shuffle: Whether the data should be shuffled.
hparams: A namedtuple of hyperparameters. This function expects
.dropout_rate: The fraction of the input units to drop, used in dropout
layer.
use_augmentation: Use data augmentation for preprocessing.
"""
super(ImageClassifier,
self).__init__(model_export_format, model_spec, index_to_label,
num_classes, shuffle, hparams.do_fine_tuning)
self.hparams = hparams
self.model = self._create_model()
self.preprocessor = image_preprocessing.Preprocessor(
self.model_spec.input_image_shape,
num_classes,
self.model_spec.mean_rgb,
self.model_spec.stddev_rgb,
use_augmentation=use_augmentation)
def __init__(self,
data,
model_export_format,
model_spec,
shuffle=True,
train_whole_model=False,
validation_ratio=0.1,
test_ratio=0.1,
hparams=lib.get_default_hparams()):
"""Init function for ImageClassifier class.
Including splitting the raw input data into train/eval/test sets and
selecting the exact NN model to be used.
Args:
data: Raw data that could be splitted for training / validation / testing.
model_export_format: Model export format such as saved_model / tflite.
model_spec: Specification for the model.
shuffle: Whether the data should be shuffled.
train_whole_model: If true, the Hub module is trained together with the
classification layer on top. Otherwise, only train the top
classification layer.
validation_ratio: The ratio of validation data to be splitted.
test_ratio: The ratio of test data to be splitted.
hparams: A namedtuple of hyperparameters. This function expects
.dropout_rate: The fraction of the input units to drop, used in dropout
layer.
"""
super(ImageClassifier,
self).__init__(data, model_export_format, model_spec, shuffle,
train_whole_model, validation_ratio, test_ratio)
# Gets pre_trained models.
if model_export_format != mef.ModelExportFormat.TFLITE:
raise ValueError('Model export mode %s is not supported currently.' %
str(model_export_format))
self.pre_trained_model_spec = model_spec
# Generates training, validation and testing data.
if validation_ratio + test_ratio >= 1.0:
raise ValueError(
'The total ratio for validation and test data should be less than 1.0.'
)
self.validation_data, rest_data = data.split(
validation_ratio, shuffle=shuffle)
self.test_data, self.train_data = rest_data.split(
test_ratio, shuffle=shuffle)
# Checks dataset parameter.
if self.train_data.size == 0:
raise ValueError('Training dataset is empty.')
# Creates the classifier model for retraining.
module_layer = hub.KerasLayer(
self.pre_trained_model_spec.uri, trainable=train_whole_model)
self.model = lib.build_model(module_layer, hparams,
self.pre_trained_model_spec.input_image_shape,
data.num_classes)
def get_default_hparams():
"""Returns a fresh HParams object initialized to default values."""
default_hub_hparams = hub_lib.get_default_hparams()
as_dict = default_hub_hparams._asdict()
as_dict.update(
train_epochs=10,
do_fine_tuning=False,
batch_size=64,
learning_rate=0.004,
dropout_rate=0.2,
warmup_steps=None,
model_dir=tempfile.mkdtemp(),
)
default_hparams = HParams(**as_dict)
return default_hparams
def get_hub_lib_hparams(**kwargs):
"""Gets the hyperparameters for the tensorflow hub's library."""
hparams = hub_lib.get_default_hparams()
return train_image_classifier_lib.add_params(hparams, **kwargs)
def create(train_data,
model_export_format=mef.ModelExportFormat.TFLITE,
model_spec=ms.mobilenet_v2_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
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tempfile
from absl import app
from absl import flags
from absl import logging
import six
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow_hub.tools.make_image_classifier import make_image_classifier_lib as lib
_DEFAULT_HPARAMS = lib.get_default_hparams()
flags.DEFINE_string(
"image_dir", None,
"A directory with subdirectories of images, one per class. "
"If unset, the TensorFlow Flowers example dataset will be used. "
"Internally, the dataset is split into training and validation pieces.")
flags.DEFINE_string(
"tfhub_module",
"https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4",
"Which TF Hub module to use. Must be a module in TF2/SavedModel format "
"for computing image feature vectors.")
flags.DEFINE_integer(
"image_size", None,
"The height and width of images to feed into --tfhub_module. "
"(For now, must be set manually for modules with variable input size.)")
