diff --git a/docs/common_saved_model_apis/images.md b/docs/common_saved_model_apis/images.md index 9a01015c6..2b2cc217d 100644 --- a/docs/common_saved_model_apis/images.md +++ b/docs/common_saved_model_apis/images.md @@ -1,4 +1,4 @@ - + # Common SavedModel APIs for Image Tasks @@ -72,8 +72,6 @@ Reusable SavedModels for image feature vectors are used in * the Colab tutorial [Retraining an Image Classifier](https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/tf2_image_retraining.ipynb), -* the command-line tool - [make_image_classifier](https://github.com/tensorflow/hub/tree/master/tensorflow_hub/tools/make_image_classifier). diff --git a/tensorflow_hub/BUILD b/tensorflow_hub/BUILD index cb488dde4..ca0adbd8e 100644 --- a/tensorflow_hub/BUILD +++ b/tensorflow_hub/BUILD @@ -486,12 +486,6 @@ py_library( name = "expect_numpy_installed", ) -# We expect absl to already be installed on the system, e.g. via -# `pip install absl-py` -py_library( - name = "expect_absl_py_installed", -) - # We expect the runtime libraries for generated Python protocol messages # to already be installed on the system, e.g. via # `pip install protobuf` diff --git a/tensorflow_hub/pip_package/BUILD b/tensorflow_hub/pip_package/BUILD index 78ea406b6..d7f3db453 100644 --- a/tensorflow_hub/pip_package/BUILD +++ b/tensorflow_hub/pip_package/BUILD @@ -25,6 +25,5 @@ sh_binary( srcs = ["build_pip_package.sh"], data = [ "//tensorflow_hub", - "//tensorflow_hub/tools/make_image_classifier", ], ) diff --git a/tensorflow_hub/pip_package/setup.py b/tensorflow_hub/pip_package/setup.py index a776e62d5..c6d1409f0 100644 --- a/tensorflow_hub/pip_package/setup.py +++ b/tensorflow_hub/pip_package/setup.py @@ -46,25 +46,19 @@ setup( name=project_name, # Automatic: tensorflow_hub, etc. Case insensitive. version=version.replace('-', ''), - description=('TensorFlow Hub is a library to foster the publication, ' - 'discovery, and consumption of reusable parts of machine ' - 'learning models.'), + description=( + 'TensorFlow Hub is a library to foster the publication, ' + 'discovery, and consumption of reusable parts of machine ' + 'learning models.' + ), long_description='', url='https://github.com/tensorflow/hub', author='Google LLC', author_email='packages@tensorflow.org', packages=find_packages(), install_requires=REQUIRED_PACKAGES, - extras_require={ - 'make_image_classifier': ['keras_preprocessing[image]'], - }, - entry_points={ - 'console_scripts': [ - ('make_image_classifier = ' - 'tensorflow_hub.tools.make_image_classifier.' - 'make_image_classifier:run_main [make_image_classifier]'), - ], - }, + extras_require={}, + entry_points={}, # PyPI package information. classifiers=[ 'Development Status :: 4 - Beta', @@ -86,6 +80,8 @@ 'Topic :: Software Development :: Libraries :: Python Modules', ], license='Apache 2.0', - keywords=('tensorflow machine learning share module subgraph component hub ' - 'embedding retraining transfer'), + keywords=( + 'tensorflow machine learning share module subgraph component hub ' + 'embedding retraining transfer' + ), ) diff --git a/tensorflow_hub/tools/make_image_classifier/BUILD b/tensorflow_hub/tools/make_image_classifier/BUILD deleted file mode 100644 index 67ccf9dab..000000000 --- a/tensorflow_hub/tools/make_image_classifier/BUILD +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved. -# -# 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. -# ============================================================================== -licenses(["notice"]) - -package( - default_applicable_licenses = ["//tensorflow_hub:license"], - default_visibility = [ - "//:__subpackages__", - "//tensorflow_hub:__subpackages__", - ], -) - -# A library with reusable pieces of make_image_classifier. -py_library( - name = "make_image_classifier_lib", - srcs = ["make_image_classifier_lib.py"], - srcs_version = "PY3", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow_hub:expect_tensorflow_installed", - "//tensorflow_hub", - # "tensorflow/lite/python:analyzer", - # "tensorflow/lite/python:lite", - # "tensorflow/lite/python/authoring", - ], -) - -# The make_image_classifier script as a py_library (private, for testing only). -py_library( - name = "make_image_classifier_main", - srcs = ["make_image_classifier.py"], - srcs_version = "PY3", - deps = [ - ":make_image_classifier_lib", - "//tensorflow_hub:expect_tensorflow_installed", - "//tensorflow_hub", - ], -) - -# The make_image_classifier script as a py_binary. -py_binary( - name = "make_image_classifier", - srcs = ["make_image_classifier.py"], - python_version = "PY3", - srcs_version = "PY3", - deps = [ - ":make_image_classifier_main", - ], -) - -py_test( - name = "make_image_classifier_test", - srcs = ["make_image_classifier_test.py"], - python_version = "PY3", - deps = [ - ":make_image_classifier_lib", - ":make_image_classifier_main", - "//tensorflow_hub:expect_absl_py_installed", # "/testing:flagsaver" - "//tensorflow_hub:expect_tensorflow_installed", - "//tensorflow_hub", - ], -) diff --git a/tensorflow_hub/tools/make_image_classifier/README.md b/tensorflow_hub/tools/make_image_classifier/README.md deleted file mode 100644 index 064515bca..000000000 --- a/tensorflow_hub/tools/make_image_classifier/README.md +++ /dev/null @@ -1,148 +0,0 @@ -# Making your own TensorFlow model for image classification - -The `make_image_classifier` tool comes with the tensorflow_hub library -and lets you build and train a TensorFlow model for image classification -from the command line, no coding required. The tool needs -a number of example images for each class (many dozens or hundreds), -but a default ("TF Flowers") is provided. - -If you are a developer looking for a coding example, please take a look at -[examples/colab/tf2_image_retraining.ipynb](https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/tf2_image_retraining.ipynb) -which demonstrates the key techniques of this program in your browser. - - -## Installation - -This tool requires tensorflow and tensorflow_hub libraries, -which can be installed with: - -```shell -$ pip install "tensorflow~=2.0" -$ pip install "tensorflow-hub[make_image_classifier]~=0.6" -``` - -After installation, the `make_image_classifier` executable is available -on the command line: - -```shell -$ make_image_classifier --help -``` - -This tool tends to run much faster with a GPU, if TensorFlow is installed -to use it. To do so, you need to install GPU drivers per -[tensorflow.org/install/gpu](https://www.tensorflow.org/install/gpu) -and use pip package `"tensorflow-gpu~=2.0"`. - - -## Basic Usage - -Basic usage of the tool looks like - -```shell -$ make_image_classifier \ - --image_dir my_image_dir \ - --tfhub_module https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4 \ - --image_size 224 \ - --saved_model_dir my_dir/new_model \ - --labels_output_file class_labels.txt \ - --tflite_output_file new_mobile_model.tflite \ - --summaries_dir my_log_dir -``` - -The `--image_dir` is a directory of subdirectories of images, defining -the classes you want your model to distinguish. Say you wanted to -classify your photos of pets to be cat, dog or rabbit. Then you would -arrange JPEG files of such photos in a directory structure like - -``` -my_image_dir -|-- cat -| |-- a_feline_photo.jpg -| |-- another_cat_pic.jpg -| `-- ... -|-- dog -| |-- PuppyInBasket.JPG -| |-- walking_the_dog.jpeg -| `-- ... -`-- rabbit - |-- IMG87654321.JPG - |-- my_fluffy_rabbit.JPEG - `-- ... -``` - -Good training results need many images (many dozens, possibly hundreds -per class). - -**Note:** For a quick demo, omit --image_dir. This will download and use -the "TF Flowers" dataset and train a model to classify photos of flowers -as daisy, dandelion, rose, sunflower or tulip. - -The `--tfhub_module` is the URL of a pre-trained model piece, or "module", -on [TensorFlow Hub](https://tfhub.dev). You can point your browser to the -module URL to see documentation for it. This tool requires a module -for image feature extraction in TF2 format. You can find them on TF Hub with -[this search](https://tfhub.dev/s?module-type=image-feature-vector&q=tf2). - -Images are resized to the given `--image_size` after reading from -disk. It depends on the TF Hub module whether it accepts only a fixed size -(in which case you can omit this flag) or an arbitrary size (in which -case you should start off by setting this to the standard value -advertised in the module documentation). - -Model training consumes your input data multiple times ("epochs"). -Some part of the data is set aside as validation data; the partially -trained model is evaluated on that after each epoch. You can see -progress bars and accuracy indicators on the console. - -After training, the given `--saved_model_dir` is created and filled -with several files that represent the complete image classification model -in TensorFlow's SavedModel format. This can be deployed to TensorFlow Serving. - -If `--labels_output_file` is given, the names of the classes are written -to that text file, one per line, in the same order as they appear -in the predictions output by the model. - -If `--tflite_output_file` is given, the complete image classification model -is written to that file in TensorFlow Lite's model format ("flatbuffers"). -This can be deployed to TF Lite on mobile devices. -If you are not deploying to TF Lite, you can simply omit this flag. - -If `--summaries_dir` is given, you can monitor your model training -on TensorBoard. See [this guide](https://www.tensorflow.org/tensorboard/get_started) -on how to enable TensorBoard. - -If you set all the flags as in the example above, you can test the -resulting TF Lite model with -[tensorflow/lite/examples/python/label_image.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/python/label_image.py) -by downloading that program and running on an image like - -```shell -python label_image.py \ - --input_mean 0 --input_std 255 \ - --model_file new_mobile_model.tflite --label_file class_labels.txt \ - --image my_image_dir/cat/a_feline_photo.jpg # <<< Adjust filename. -``` - - -## Advanced usage - -Additional command-line flags let you control the training process. -In particular, you can increase `--train_epochs` to train more, -and set the `--learning_rate` and `--momentum` for the SGD optimizer. - -Also, you can set `--do_fine_tuning` to train the TensorFlow Hub -module together with the classifier. - -There is other hyperparameters for regularization such as -`--l1_regularizer` and `--l2_regularizer`, and for data augmentations -such as `--rotation_range` and `--horizontal_flip`. Generally, the -default values can give a good performance. You can find a full list -of hyperparameters available in `make_image_classifier.py` and their -default values in `make_image_classifier_lib.py`. - -With `tensorflow>=2.5` and `tensorflow-hub>=0.12`, you can control whether to -read input with a tf.data.Dataset and use TF ops for preprocessing using the -`use_tf_data_input` flag. Note that the shear data augmentation is not -supported in this mode. If set to `False`, Keras' legacy Python -ImageDataGenerator with numpy ops will be used for data augmentation and other -preprocessing. diff --git a/tensorflow_hub/tools/make_image_classifier/__init__.py b/tensorflow_hub/tools/make_image_classifier/__init__.py deleted file mode 100644 index b9863d00d..000000000 --- a/tensorflow_hub/tools/make_image_classifier/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved. -# -# 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. -# ============================================================================== diff --git a/tensorflow_hub/tools/make_image_classifier/make_image_classifier.py b/tensorflow_hub/tools/make_image_classifier/make_image_classifier.py deleted file mode 100644 index 4122335b2..000000000 --- a/tensorflow_hub/tools/make_image_classifier/make_image_classifier.py +++ /dev/null @@ -1,298 +0,0 @@ -# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved. -# -# 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. -# ============================================================================== -"""Trains a TensorFlow model based on directories of images. - -This program builds, trains and exports a TensorFlow 2.x model that classifies -natural images (photos) into a fixed set of classes. The classes are learned -from a user-supplied dataset of images, stored as a directory of subdirectories -of JPEG images, each subdirectory representing one class. - -The model is built from a pre-trained image feature vector module from -TensorFlow Hub (in its TF2/SavedModel format, not the legacy TF1 Hub format) -followed by a linear classifier. The linear classifier, and optionally also -the TF Hub module, are trained on the new dataset. TF Hub offers a variety of -suitable modules with various size/accuracy tradeoffs. - -The resulting model can be exported in TensorFlow's standard SavedModel format -and as a .tflite file for deployment to mobile devices with TensorFlow Lite. -TODO(b/139467904): Add support for post-training model optimization. - -For more information, please see the README file next to the source code, -https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/README.md -""" - -# NOTE: This is an expanded, command-line version of -# https://github.com/tensorflow/hub/blob/master/examples/colab/tf2_image_retraining.ipynb -# PLEASE KEEP THEM IN SYNC, such that running tests for this program -# provides assurance that the code in the colab notebook works. - -import io -import tempfile - -from absl import app -from absl import flags -from absl import logging -from distutils.version import LooseVersion -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.)") -flags.DEFINE_string( - "saved_model_dir", None, - "The final model is exported as a SavedModel directory with this name.") -flags.DEFINE_string( - "tflite_output_file", None, - "The final model is exported as a .tflite flatbuffers file with this name.") -flags.DEFINE_string( - "labels_output_file", None, - "Where to save the labels (that is, names of image subdirectories). " - "The lines in this file appear in the same order as the predictions " - "of the model.") -flags.DEFINE_string( - "summaries_dir", None, - "Where to save summary logs for TensorBoard.") -flags.DEFINE_float( - "assert_accuracy_at_least", None, - "If set, the program fails if the validation accuracy at the end of " - "training is less than this number (between 0 and 1), and no export of " - "the trained model happens.") -flags.DEFINE_integer( - "train_epochs", _DEFAULT_HPARAMS.train_epochs, - "Training will do this many iterations over the dataset.") -flags.DEFINE_bool( - "do_fine_tuning", _DEFAULT_HPARAMS.do_fine_tuning, - "If set, the --tfhub_module is trained together with the rest of " - "the model being built.") -flags.DEFINE_integer( - "batch_size", _DEFAULT_HPARAMS.batch_size, - "Each training step samples a batch of this many images " - "from the training data. (You may need to shrink this when using a GPU " - "and getting out-of-memory errors. Avoid values below 8 when re-training " - "modules that use batch normalization.)") -flags.DEFINE_float( - "learning_rate", _DEFAULT_HPARAMS.learning_rate, - "The learning rate to use for gradient descent training.") -flags.DEFINE_float( - "momentum", _DEFAULT_HPARAMS.momentum, - "The momentum parameter to use for gradient descent training.") -flags.DEFINE_float( - "dropout_rate", _DEFAULT_HPARAMS.dropout_rate, - "The fraction of the input units to drop, used in dropout layer.") -flags.DEFINE_bool( - "set_memory_growth", False, - "If flag is set, memory growth functionality flag will be set as true for " - "all GPUs prior to training. " - "More details: " - "https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth" -) -flags.DEFINE_float( - "l1_regularizer", _DEFAULT_HPARAMS.l1_regularizer, - "Coefficient of L1 regularization applied on model weights.") -flags.DEFINE_float( - "l2_regularizer", _DEFAULT_HPARAMS.l2_regularizer, - "Coefficient of L2 regularization applied on model weights.") -flags.DEFINE_float("label_smoothing", _DEFAULT_HPARAMS.label_smoothing, - "Coefficient of label smoothing used in loss function.") -flags.DEFINE_float("validation_split", _DEFAULT_HPARAMS.validation_split, - "The fraction of the dataset split into a validation set") -flags.DEFINE_bool( - "do_data_augmentation", _DEFAULT_HPARAMS.do_data_augmentation, - "Whether do data augmentation on training set." - "Can use default augmentation params or specifying them.") -flags.DEFINE_integer("rotation_range", _DEFAULT_HPARAMS.rotation_range, - "Degree range for random rotation.") -flags.DEFINE_bool("horizontal_flip", _DEFAULT_HPARAMS.horizontal_flip, - "Horizontally flip images.") -flags.DEFINE_float( - "width_shift_range", _DEFAULT_HPARAMS.width_shift_range, - "Shift images horizontally by pixels(if >=1) or by ratio(if <1).") -flags.DEFINE_float( - "height_shift_range", _DEFAULT_HPARAMS.height_shift_range, - "Shift images vertically by pixels(if >=1) or by ratio(if <1).") -flags.DEFINE_float( - "shear_range", _DEFAULT_HPARAMS.shear_range, - "Shear angle in counter-clockwise direction in degrees." - "DEPRECATED: Image shear is not available if using TF 2.5 or higher.") -flags.DEFINE_float("zoom_range", _DEFAULT_HPARAMS.zoom_range, - "Range for random zoom.") -flags.DEFINE_enum("distribution_strategy", None, ["", "mirrored"], - "The distribution strategy the classifier should use.") -flags.DEFINE_bool( - "use_tf_data_input", None, - "Whether to read input with a tf.data.Dataset and use TF ops for " - "preprocessing. Use of Keras preprocessing layers is supported for TF " - "versions 2.5 or higher. If false, uses Keras' legacy Python " - "ImageDataGenerator with numpy ops.") -FLAGS = flags.FLAGS - - -def _get_hparams_from_flags(): - """Creates dict of hyperparameters from flags.""" - return lib.HParams( - train_epochs=FLAGS.train_epochs, - do_fine_tuning=FLAGS.do_fine_tuning, - batch_size=FLAGS.batch_size, - learning_rate=FLAGS.learning_rate, - momentum=FLAGS.momentum, - dropout_rate=FLAGS.dropout_rate, - l1_regularizer=FLAGS.l1_regularizer, - l2_regularizer=FLAGS.l2_regularizer, - label_smoothing=FLAGS.label_smoothing, - validation_split=FLAGS.validation_split, - do_data_augmentation=FLAGS.do_data_augmentation, - rotation_range=FLAGS.rotation_range, - horizontal_flip=FLAGS.horizontal_flip, - width_shift_range=FLAGS.width_shift_range, - height_shift_range=FLAGS.height_shift_range, - shear_range=FLAGS.shear_range, - zoom_range=FLAGS.zoom_range) - - -def _check_keras_dependencies(): - """Checks dependencies of tf.keras.preprocessing.image are present. - - This function may come to depend on flag values that determine the kind - of preprocessing being done. - - Raises: - ImportError: If dependencies are missing. - """ - try: - tf.keras.preprocessing.image.load_img(io.BytesIO()) - except ImportError: - print("\n*** Unsatisfied dependencies of keras_preprocessing.image. ***\n" - "To install them, use your system's equivalent of\n" - "pip install tensorflow_hub[make_image_classifier]\n") - raise - except Exception as e: # pylint: disable=broad-except - # Loading from dummy content as above is expected to fail in other ways. - pass - - -def _assert_accuracy(train_result, assert_accuracy_at_least): - # Fun fact: With TF1 behavior, the key was called "val_acc". - val_accuracy = train_result.history["val_accuracy"][-1] - accuracy_message = "found {:f}, expected at least {:f}".format( - val_accuracy, assert_accuracy_at_least) - if val_accuracy >= assert_accuracy_at_least: - print("ACCURACY PASSED:", accuracy_message) - else: - raise AssertionError("ACCURACY FAILED:", accuracy_message) - - -def _set_gpu_memory_growth(): - # Original code reference found here: - # https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth - gpus = tf.config.experimental.list_physical_devices("GPU") - if gpus: - # Currently, memory growth needs to be the same across GPUs - for gpu in gpus: - tf.config.experimental.set_memory_growth(gpu, True) - print("All GPUs will scale memory steadily") - else: - print("No GPUs found for set_memory_growth") - - -def main(args): - """Main function to be called by absl.app.run() after flag parsing.""" - del args - _check_keras_dependencies() - hparams = _get_hparams_from_flags() - - image_dir = FLAGS.image_dir or lib.get_default_image_dir() - - if FLAGS.set_memory_growth: - _set_gpu_memory_growth() - - use_tf_data_input = FLAGS.use_tf_data_input - # For tensorflow<2.5 TF preprocessing layers do not support distribution - # strategy. so default use_tf_data_input to True for TF >= 2.5. - if use_tf_data_input is True and (LooseVersion(tf.__version__) < - LooseVersion("2.5.0")): - raise ValueError("use_tf_data_input is not supported for tensorflow<2.5") - # For tensorflow>=2.5 default to using tf.data.Dataset and TF preprocessing - # layers. - if use_tf_data_input is None and (LooseVersion(tf.__version__) >= - LooseVersion("2.5.0")): - use_tf_data_input = True - - model, labels, train_result = lib.make_image_classifier( - FLAGS.tfhub_module, image_dir, hparams, - lib.get_distribution_strategy(FLAGS.distribution_strategy), - FLAGS.image_size, FLAGS.summaries_dir, use_tf_data_input) - if FLAGS.assert_accuracy_at_least: - _assert_accuracy(train_result, FLAGS.assert_accuracy_at_least) - print("Done with training.") - - if FLAGS.labels_output_file: - with tf.io.gfile.GFile(FLAGS.labels_output_file, "w") as f: - f.write("\n".join(labels + ("",))) - print("Labels written to", FLAGS.labels_output_file) - - saved_model_dir = FLAGS.saved_model_dir - if FLAGS.tflite_output_file and not saved_model_dir: - # We need a SavedModel for conversion, even if the user did not request it. - saved_model_dir = tempfile.mkdtemp() - if saved_model_dir: - tf.saved_model.save(model, saved_model_dir) - print("SavedModel model exported to", saved_model_dir) - - if FLAGS.tflite_output_file: - converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) - lite_model_content = converter.convert() - with tf.io.gfile.GFile(FLAGS.tflite_output_file, "wb") as f: - f.write(lite_model_content) - print("TFLite model exported to", FLAGS.tflite_output_file) - - -def _ensure_tf2(): - """Ensure running with TensorFlow 2 behavior. - - This function is safe to call even before flags have been parsed. - - Raises: - ImportError: If tensorflow is too old for proper TF2 behavior. - """ - logging.info("Running with tensorflow %s and hub %s", - tf.__version__, hub.__version__) - if not tf.executing_eagerly(): - raise ImportError("Sorry, this program needs TensorFlow 2.") - - -def run_main(): - """Entry point equivalent to executing this file.""" - _ensure_tf2() - app.run(main) - - -if __name__ == "__main__": - run_main() diff --git a/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py b/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py deleted file mode 100644 index 4d229429b..000000000 --- a/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py +++ /dev/null @@ -1,412 +0,0 @@ -# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved. -# -# 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. -# ============================================================================== -"""Trains a TensorFlow model based on directories of images. - -This library provides the major pieces for make_image_classifier (see there). -""" - -import collections -import contextlib - -import tensorflow as tf -import tensorflow_hub as hub - -_DEFAULT_IMAGE_URL = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz" - - -def get_default_image_dir(): - """Returns the path to a default image dataset, downloading it if needed.""" - return tf.keras.utils.get_file("flower_photos", - _DEFAULT_IMAGE_URL, untar=True) - - -class NoStrategy: - scope = contextlib.contextmanager(lambda _: iter(range(1))) - - -def get_distribution_strategy(distribution_strategy_name): - if distribution_strategy_name == "mirrored": - return tf.distribute.MirroredStrategy() - elif not distribution_strategy_name: - return NoStrategy() - else: - raise ValueError( - "Unknown distribution strategy {}".format(distribution_strategy_name)) - - -class HParams( - collections.namedtuple("HParams", [ - "train_epochs", "do_fine_tuning", "batch_size", "learning_rate", - "momentum", "dropout_rate", "l1_regularizer", "l2_regularizer", - "label_smoothing", "validation_split", "do_data_augmentation", - "rotation_range", "horizontal_flip", "width_shift_range", - "height_shift_range", "shear_range", "zoom_range" - ])): - """The hyperparameters for make_image_classifier. - - train_epochs: Training will do this many iterations over the dataset. - do_fine_tuning: If true, the Hub module is trained together with the - classification layer on top. - batch_size: Each training step samples a batch of this many images. - learning_rate: The learning rate to use for gradient descent training. - momentum: The momentum parameter to use for gradient descent training. - dropout_rate: The fraction of the input units to drop, used in dropout layer. - """ - - -def get_default_hparams(): - """Returns a fresh HParams object initialized to default values.""" - return HParams( - train_epochs=5, - do_fine_tuning=False, - batch_size=32, - learning_rate=0.005, - momentum=0.9, - dropout_rate=0.2, - l1_regularizer=0.0, - l2_regularizer=0.0001, - label_smoothing=0.1, - validation_split=0.2, - do_data_augmentation=False, - rotation_range=40, - horizontal_flip=True, - width_shift_range=0.2, - height_shift_range=0.2, - shear_range=0, - zoom_range=0.2) - - -def _get_data_with_keras(image_dir, image_size, batch_size, validation_split, - do_data_augmentation, augmentation_params): - """Gets training and validation data via keras_preprocessing. - - Args: - image_dir: A Python string with the name of a directory that contains - subdirectories of images, one per class. - image_size: A list or tuple with 2 Python integers specifying - the fixed height and width to which input images are resized. - batch_size: A Python integer with the number of images per batch of - training and validation data. - validation_split: A float representing the fraction of the dataset split - into a validation set. - do_data_augmentation: An optional boolean, controlling whether the - training dataset is augmented by randomly distorting input images. - augmentation_params: A dictionary containing the augmentation params as keys - and their respective values. - - Returns: - A nested tuple ((train_data, train_size), - (valid_data, valid_size), labels) where: - train_data, valid_data: Generators for use with Model.fit_generator, - each yielding tuples (images, labels) where - images is a float32 Tensor of shape [batch_size, height, width, 3] - with pixel values in range [0,1], - labels is a float32 Tensor of shape [batch_size, num_classes] - with one-hot encoded classes. - train_size, valid_size: Python integers with the numbers of training - and validation examples, respectively. - labels: A tuple of strings with the class labels (subdirectory names). - The index of a label in this tuple is the numeric class id. - """ - datagen_kwargs = dict(rescale=1. / 255, validation_split=validation_split) - dataflow_kwargs = dict(target_size=image_size, batch_size=batch_size, - interpolation="bilinear") - - valid_datagen = tf.keras.preprocessing.image.ImageDataGenerator( - **datagen_kwargs) - valid_generator = valid_datagen.flow_from_directory( - image_dir, subset="validation", shuffle=False, **dataflow_kwargs) - - if do_data_augmentation and len(augmentation_params): - datagen_kwargs.update(**augmentation_params) - train_datagen = tf.keras.preprocessing.image.ImageDataGenerator( - **datagen_kwargs) - else: - train_datagen = valid_datagen - train_generator = train_datagen.flow_from_directory( - image_dir, subset="training", shuffle=True, **dataflow_kwargs) - - indexed_labels = [(index, label) - for label, index in train_generator.class_indices.items()] - sorted_indices, sorted_labels = zip(*sorted(indexed_labels)) - assert sorted_indices == tuple(range(len(sorted_labels))) - return ((train_generator, train_generator.samples), - (valid_generator, valid_generator.samples), - sorted_labels) - - -def _get_data_as_datasets(image_dir, image_size, hparams): - """Gets training and validation data via tf.data.Dataset. - - Args: - image_dir: A Python string with the name of a directory that contains - subdirectories of images, one per class. - image_size: A list or tuple with 2 Python integers specifying the fixed - height and width to which input images are resized. - hparams: A HParams object with hyperparameters controlling the training. - - Returns: - A nested tuple ((train_data, train_size), - (valid_data, valid_size), labels) where: - train_data, valid_data: tf.data.Dataset for use with Model.fit, each - yielding batch of tuples (images, labels) where - images is a float32 Tensor of shape [batch_size, height, width, 3] - with pixel values in range [0,1], - labels is a float32 Tensor of shape [batch_size, num_classes] - with one-hot encoded classes. - train_size, valid_size: Python integers with the numbers of training - and validation examples, respectively. - labels: A tuple of strings with the class labels (subdirectory names). - The index of a label in this tuple is the numeric class id. - """ - # Check if hparam.shear_range is set. If yes, throw an error since shear is - # not supported when using preprocessing layers. - if hparams.shear_range != 0: - raise ValueError("Found non-zero value for shear_range. Shear is not " - "supported when using reading input with tf.data.Dataset " - "and using preprocessing layers.") - - train_ds = tf.keras.preprocessing.image_dataset_from_directory( - image_dir, - validation_split=hparams.validation_split, - subset="training", - label_mode="categorical", - # Seed needs to provided when using validation_split and shuffle = True. - # A fixed seed is used so that the validation set is stable across runs. - seed=123, - image_size=image_size, - batch_size=1) - class_names = tuple(train_ds.class_names) - train_size = train_ds.cardinality().numpy() - train_ds = train_ds.unbatch().batch(hparams.batch_size) - train_ds = train_ds.repeat() - - normalization_layer = tf.keras.layers.experimental.preprocessing.Rescaling( - 1. / 255) - preprocessing_model = tf.keras.Sequential([normalization_layer]) - if hparams.do_data_augmentation: - preprocessing_model.add( - tf.keras.layers.experimental.preprocessing.RandomRotation( - hparams.rotation_range)) - preprocessing_model.add( - tf.keras.layers.experimental.preprocessing.RandomTranslation( - 0, hparams.width_shift_range)) - preprocessing_model.add( - tf.keras.layers.experimental.preprocessing.RandomTranslation( - hparams.height_shift_range, 0)) - # Like the old tf.keras.preprocessing.image.ImageDataGenerator(), - # image sizes are fixed when reading, and then a random zoom is applied. - # If all training inputs are larger than image_size, one could also use - # RandomCrop with a batch size of 1 and rebatch later. - preprocessing_model.add( - tf.keras.layers.experimental.preprocessing.RandomZoom( - hparams.zoom_range, hparams.zoom_range)) - if hparams.horizontal_flip: - preprocessing_model.add( - tf.keras.layers.experimental.preprocessing.RandomFlip( - mode="horizontal")) - train_ds = train_ds.map(lambda images, labels: - (preprocessing_model(images), labels)) - - val_ds = tf.keras.preprocessing.image_dataset_from_directory( - image_dir, - validation_split=hparams.validation_split, - subset="validation", - label_mode="categorical", - seed=123, - shuffle=False, - image_size=image_size, - batch_size=1) - valid_size = val_ds.cardinality().numpy() - val_ds = val_ds.unbatch().batch(hparams.batch_size) - val_ds = val_ds.map(lambda images, labels: - (normalization_layer(images), labels)) - - return ((train_ds, train_size), (val_ds, valid_size), class_names) - - -def _image_size_for_module(module_layer, requested_image_size=None): - """Returns the input image size to use with the given module. - - Args: - module_layer: A hub.KerasLayer initialized from a Hub module expecting - image input. - requested_image_size: An optional Python integer with the user-requested - height and width of the input image; or None. - - Returns: - A tuple (height, width) of Python integers that can be used as input - image size for the given module_layer. - - Raises: - ValueError: If requested_image_size is set but incompatible with the module. - ValueError: If the module does not specify a particular input size and - requested_image_size is not set. - """ - # TODO(b/139530454): Use a library helper function once available. - # The stop-gap code below assumes any concrete function backing the - # module call will accept a batch of images with the one accepted size. - module_image_size = tuple( - module_layer._func.__call__ # pylint:disable=protected-access - .concrete_functions[0].structured_input_signature[0][0].shape[1:3]) - if requested_image_size is None: - if None in module_image_size: - raise ValueError("Must specify an image size because " - "the selected TF Hub module specifies none.") - else: - return module_image_size - else: - requested_image_size = tf.TensorShape( - [requested_image_size, requested_image_size]) - assert requested_image_size.is_fully_defined() - if requested_image_size.is_compatible_with(module_image_size): - return tuple(requested_image_size.as_list()) - else: - raise ValueError("The selected TF Hub module expects image size {}, " - "but size {} is requested".format( - module_image_size, - tuple(requested_image_size.as_list()))) - - -def build_model(module_layer, hparams, image_size, num_classes): - """Builds the full classifier model from the given module_layer. - - If using a DistributionStrategy, call this under its `.scope()`. - Args: - module_layer: Pre-trained tfhub model layer. - hparams: A namedtuple of hyperparameters. This function expects - .dropout_rate: The fraction of the input units to drop, used in dropout - layer. - image_size: The input image size to use with the given module layer. - num_classes: Number of the classes to be predicted. - - Returns: - The full classifier model. - """ - model = tf.keras.Sequential([ - tf.keras.Input(shape=(image_size[0], image_size[1], 3)), module_layer, - tf.keras.layers.Dropout(rate=hparams.dropout_rate), - tf.keras.layers.Dense( - num_classes, - activation="softmax", - kernel_regularizer=tf.keras.regularizers.l1_l2( - l1=hparams.l1_regularizer, l2=hparams.l2_regularizer)) - ]) - print(model.summary()) - return model - - -def train_model(model, - hparams, - train_data_and_size, - valid_data_and_size, - log_dir=None): - """Trains model with the given data and hyperparameters. - - If using a DistributionStrategy, call this under its `.scope()`. - Args: - model: The tf.keras.Model from _build_model(). - 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, the number of examples returned by each - call to the generators. - train_data_and_size: A (data, size) tuple in which data is training data to - be fed in tf.keras.Model.fit(), size is a Python integer with the - numbers of training. - valid_data_and_size: A (data, size) tuple in which data is validation data - to be fed in tf.keras.Model.fit(), size is a Python integer with the - numbers of validation. - log_dir: A directory to write logs for TensorBoard into (defaults to None, - no logs will then be written). - - Returns: - The tf.keras.callbacks.History object returned by tf.keras.Model.fit(). - """ - train_data, train_size = train_data_and_size - valid_data, valid_size = valid_data_and_size - loss = tf.keras.losses.CategoricalCrossentropy( - label_smoothing=hparams.label_smoothing) - model.compile( - optimizer=tf.keras.optimizers.SGD( - learning_rate=hparams.learning_rate, momentum=hparams.momentum), - loss=loss, - metrics=["accuracy"]) - steps_per_epoch = train_size // hparams.batch_size - validation_steps = valid_size // hparams.batch_size - callbacks = [] - if log_dir != None: - callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=log_dir, - histogram_freq=1)) - return model.fit( - train_data, - epochs=hparams.train_epochs, - steps_per_epoch=steps_per_epoch, - validation_data=valid_data, - validation_steps=validation_steps, - callbacks=callbacks) - - -def make_image_classifier(tfhub_module, - image_dir, - hparams, - distribution_strategy=None, - requested_image_size=None, - log_dir=None, - use_tf_data_input=False): - """Builds and trains a TensorFLow model for image classification. - - Args: - tfhub_module: A Python string with the handle of the Hub module. - image_dir: A Python string naming a directory with subdirectories of images, - one per class. - hparams: A HParams object with hyperparameters controlling the training. - distribution_strategy: The DistributionStrategy make_image_classifier is - running with. - requested_image_size: A Python integer controlling the size of images to - feed into the Hub module. If the module has a fixed input size, this must - be omitted or set to that same value. - log_dir: A directory to write logs for TensorBoard into (defaults to None, - no logs will then be written). - use_tf_data_input: Whether to read input with a tf.data.Dataset and use TF - ops for preprocessing. - """ - augmentation_params = dict( - rotation_range=hparams.rotation_range, - horizontal_flip=hparams.horizontal_flip, - width_shift_range=hparams.width_shift_range, - height_shift_range=hparams.height_shift_range, - shear_range=hparams.shear_range, - zoom_range=hparams.zoom_range) - - with distribution_strategy.scope(): - module_layer = hub.KerasLayer( - tfhub_module, trainable=hparams.do_fine_tuning) - image_size = _image_size_for_module(module_layer, requested_image_size) - print("Using module {} with image size {}".format(tfhub_module, image_size)) - if use_tf_data_input: - train_data_and_size, valid_data_and_size, labels = _get_data_as_datasets( - image_dir, image_size, hparams) - else: - train_data_and_size, valid_data_and_size, labels = _get_data_with_keras( - image_dir, image_size, hparams.batch_size, hparams.validation_split, - hparams.do_data_augmentation, augmentation_params) - print("Found", len(labels), "classes:", ", ".join(labels)) - model = build_model(module_layer, hparams, image_size, len(labels)) - train_result = train_model(model, hparams, train_data_and_size, - valid_data_and_size, log_dir) - return model, labels, train_result diff --git a/tensorflow_hub/tools/make_image_classifier/make_image_classifier_test.py b/tensorflow_hub/tools/make_image_classifier/make_image_classifier_test.py deleted file mode 100644 index 286b02f0f..000000000 --- a/tensorflow_hub/tools/make_image_classifier/make_image_classifier_test.py +++ /dev/null @@ -1,230 +0,0 @@ -# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved. -# -# 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. -# ============================================================================== -"""Unit tests for make_image_classifier_lib.py and make_image_classifier.py. - -For now, we keep a single unit test for the library and its command-line -driver, because the latter is the best way to achieve end-to-end testing. -""" - -import os -import random -import sys - -from absl import logging -from absl.testing import flagsaver -from absl.testing import parameterized -from distutils.version import LooseVersion -import numpy as np -import tensorflow as tf -import tensorflow_hub as hub - -from tensorflow_hub.tools.make_image_classifier import make_image_classifier -from tensorflow_hub.tools.make_image_classifier import make_image_classifier_lib - - -def _fill_image(rgb, image_size): - r, g, b = rgb - return np.broadcast_to(np.array([[[r, g, b]]], dtype=np.uint8), - shape=(image_size, image_size, 3)) - - -def _write_filled_jpeg_file(path, rgb, image_size): - tf.keras.preprocessing.image.save_img(path, _fill_image(rgb, image_size), - "channels_last", "jpeg") - - -class MakeImageClassifierTest(tf.test.TestCase, parameterized.TestCase): - IMAGE_SIZE = 24 - IMAGES_PER_CLASS = 20 - CMY_NAMES_AND_RGB_VALUES = (("cyan", (0, 255, 255)), - ("magenta", (255, 0, 255)), - ("yellow", (255, 255, 0))) - DEFAULT_FLAGS = dict(image_size=IMAGE_SIZE, train_epochs=10, - batch_size=8, learning_rate=0.1, momentum=0.0) - - def _write_cmy_dataset(self): - path = os.path.join(self.get_temp_dir(), "cmy_image_dir") - os.mkdir(path) # Fails if exists. - for class_name, rgb in self.CMY_NAMES_AND_RGB_VALUES: - class_subdir = os.path.join(path, class_name) - os.mkdir(class_subdir) - for i in range(self.IMAGES_PER_CLASS): - _write_filled_jpeg_file( - os.path.join(class_subdir, "img_%s_%03d.jpeg" % (class_name, i)), - rgb, self.IMAGE_SIZE) - return path - - def _write_random_dataset(self): - path = os.path.join(self.get_temp_dir(), "random_image_dir") - os.mkdir(path) # Fails if exists. - for class_name in ("ami", "baz", "zrh"): - class_subdir = os.path.join(path, class_name) - os.mkdir(class_subdir) - for i in range(self.IMAGES_PER_CLASS): - _write_filled_jpeg_file( - os.path.join(class_subdir, "img_%s_%03d.jpeg" % (class_name, i)), - [random.uniform(0, 255) for _ in range(3)], - self.IMAGE_SIZE) - return path - - def _export_global_average_model(self, has_fixed_input_size=True): - if has_fixed_input_size: - input_size = (self.IMAGE_SIZE, self.IMAGE_SIZE, 3) - dirname = "global_average_fixed_size" - else: - input_size = (None, None, 3) - dirname = "global_average_variable_size" - path = os.path.join(self.get_temp_dir(), dirname) - inputs = tf.keras.Input(input_size) - outputs = tf.keras.layers.GlobalAveragePooling2D()(inputs) - model = tf.keras.Model(inputs, outputs) - model.build((None,) + input_size) - model.save(path, save_format="tf") - return path - - def _load_labels(self, filename): - with tf.io.gfile.GFile(filename, "r") as f: - return [label.strip("\n") for label in f] - - def _load_lite_model(self, filename): - """Returns a numpy-to-numpy wrapper for the model in a .tflite file.""" - self.assertTrue(os.path.isfile(filename)) - with tf.io.gfile.GFile(filename, "rb") as f: - model_content = f.read() - interpreter = tf.lite.Interpreter(model_content=model_content) - def lite_model(images): - interpreter.allocate_tensors() - input_index = interpreter.get_input_details()[0]['index'] - interpreter.set_tensor(input_index, images) - interpreter.invoke() - output_index = interpreter.get_output_details()[0]['index'] - return interpreter.get_tensor(output_index) - return lite_model - - @parameterized.named_parameters(("WithLegacyInput", False), - ("WithTFDataInput", True)) - def testEndToEndSuccess(self, use_tf_data_input): - if use_tf_data_input and (LooseVersion(tf.__version__) < - LooseVersion("2.5.0")): - return - logging.info("Using testdata in %s", self.get_temp_dir()) - avg_model_dir = self._export_global_average_model() - image_dir = self._write_cmy_dataset() - saved_model_dir = os.path.join(self.get_temp_dir(), "final_saved_model") - saved_model_expected_file = os.path.join(saved_model_dir, "saved_model.pb") - tflite_output_file = os.path.join(self.get_temp_dir(), "final_model.tflite") - labels_output_file = os.path.join(self.get_temp_dir(), "labels.txt") - # Make sure we don't test for pre-existing files. - self.assertFalse(os.path.isfile(saved_model_expected_file)) - self.assertFalse(os.path.isfile(tflite_output_file)) - self.assertFalse(os.path.isfile(labels_output_file)) - - with flagsaver.flagsaver( - image_dir=image_dir, - tfhub_module=avg_model_dir, - # This dataset is expected to be fit perfectly. - assert_accuracy_at_least=0.9, - use_tf_data_input=use_tf_data_input, - saved_model_dir=saved_model_dir, - tflite_output_file=tflite_output_file, - labels_output_file=labels_output_file, - **self.DEFAULT_FLAGS): - make_image_classifier.main([]) - - # Test that the SavedModel was written. - self.assertTrue(os.path.isfile(saved_model_expected_file)) - - # Test that the TFLite model works. - labels = self._load_labels(labels_output_file) - lite_model = self._load_lite_model(tflite_output_file) - for class_name, rgb in self.CMY_NAMES_AND_RGB_VALUES: - input_batch = (_fill_image(rgb, self.IMAGE_SIZE)[None, ...] - / np.array(255., dtype=np.float32)) - output_batch = lite_model(input_batch) - prediction = labels[np.argmax(output_batch[0])] - self.assertEqual(class_name, prediction) - - @parameterized.named_parameters(("WithLegacyInput", False), - ("WithTFDataInput", True)) - def testEndToEndAccuracyFailure(self, use_tf_data_input): - if use_tf_data_input and (LooseVersion(tf.__version__) < - LooseVersion("2.5.0")): - return - logging.info("Using testdata in %s", self.get_temp_dir()) - avg_model_dir = self._export_global_average_model() - image_dir = self._write_random_dataset() - - with flagsaver.flagsaver( - image_dir=image_dir, - tfhub_module=avg_model_dir, - # This is expected to fail for this random dataset. - assert_accuracy_at_least=0.9, - use_tf_data_input=use_tf_data_input, - **self.DEFAULT_FLAGS): - with self.assertRaisesRegex(AssertionError, "ACCURACY FAILED"): - make_image_classifier.main([]) - - def testImageSizeForModuleWithFixedInputSize(self): - model_dir = self._export_global_average_model(has_fixed_input_size=True) - module_layer = hub.KerasLayer(model_dir) - self.assertTupleEqual( - (self.IMAGE_SIZE, self.IMAGE_SIZE), - make_image_classifier_lib._image_size_for_module(module_layer, None)) - self.assertTupleEqual( - (self.IMAGE_SIZE, self.IMAGE_SIZE), - make_image_classifier_lib._image_size_for_module(module_layer, - self.IMAGE_SIZE)) - with self.assertRaisesRegex(ValueError, "image size"): - make_image_classifier_lib._image_size_for_module( - module_layer, self.IMAGE_SIZE + 1) - - def testImageSizeForModuleWithVariableInputSize(self): - model_dir = self._export_global_average_model(has_fixed_input_size=False) - module_layer = hub.KerasLayer(model_dir) - self.assertTupleEqual( - (self.IMAGE_SIZE, self.IMAGE_SIZE), - make_image_classifier_lib._image_size_for_module(module_layer, - self.IMAGE_SIZE)) - self.assertTupleEqual( - (2 * self.IMAGE_SIZE, 2 * self.IMAGE_SIZE), - make_image_classifier_lib._image_size_for_module(module_layer, - 2 * self.IMAGE_SIZE)) - with self.assertRaisesRegex(ValueError, "none"): - make_image_classifier_lib._image_size_for_module(module_layer, None) - - def testGetDistributionStrategy(self): - self.assertIsInstance( - make_image_classifier_lib.get_distribution_strategy(None), - make_image_classifier_lib.NoStrategy) - self.assertIsInstance( - make_image_classifier_lib.get_distribution_strategy(""), - make_image_classifier_lib.NoStrategy) - - self.assertIsInstance( - make_image_classifier_lib.get_distribution_strategy("mirrored"), - tf.distribute.MirroredStrategy) - - with self.assertRaisesRegex(ValueError, - "Unknown distribution strategy other"): - make_image_classifier_lib.get_distribution_strategy("other") - - -if __name__ == "__main__": - try: - make_image_classifier._ensure_tf2() - except ImportError as e: - print("Skipping tests:", str(e)) - sys.exit(0) - tf.test.main()