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() model, labels, train_result = lib.make_image_classifier( FLAGS.tfhub_module, image_dir, hparams, FLAGS.image_size) 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 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 retrain(self, image_dir, output_graph, training_steps, learning_rate, desired_training_accuracy=100.0, desired_validation_accuracy=100.0, flip_left_right=True, random_crop=30, random_scale=30, random_brightness=30): _check_keras_dependencies() hparams = _get_hparams(train_epochs=training_steps, learning_rate=learning_rate) image_size = self.shape[0] tfhub_module = "https://tfhub.dev/google/" + self.architecture model, labels, train_result = lib.make_image_classifier( tfhub_module, image_dir, hparams, image_size) print("Done with training.") labels_output_file = output_graph + ".txt" with tf.io.gfile.GFile(labels_output_file, "w") as f: f.write("\n".join(labels + ("", ))) print("Labels written to", labels_output_file) saved_model_dir = tempfile.mkdtemp() tf.saved_model.save(model, saved_model_dir) print("SavedModel model exported to", saved_model_dir) tflite_output_file = output_graph + ".tflite" converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_LATENCY] lite_model_content = converter.convert() with tf.io.gfile.GFile(tflite_output_file, "wb") as f: f.write(lite_model_content) print("TFLite model exported to", tflite_output_file)