diff options
Diffstat (limited to 'tensorflow/examples/image_retraining/retrain.py')
-rw-r--r-- | tensorflow/examples/image_retraining/retrain.py | 82 |
1 files changed, 16 insertions, 66 deletions
diff --git a/tensorflow/examples/image_retraining/retrain.py b/tensorflow/examples/image_retraining/retrain.py index ebddfb20f4..3549891461 100644 --- a/tensorflow/examples/image_retraining/retrain.py +++ b/tensorflow/examples/image_retraining/retrain.py @@ -69,18 +69,11 @@ to validate that you have gathered good training data, but if you want to deploy on resource-limited platforms, you can try the `--architecture` flag with a Mobilenet model. For example: -Run floating-point version of mobilenet: ```bash python tensorflow/examples/image_retraining/retrain.py \ --image_dir ~/flower_photos --architecture mobilenet_1.0_224 ``` -Run quantized version of mobilenet: -```bash -python tensorflow/examples/image_retraining/retrain.py \ - --image_dir ~/flower_photos/ --architecture mobilenet_1.0_224_quantized -``` - There are 32 different Mobilenet models to choose from, with a variety of file size and latency options. The first number can be '1.0', '0.75', '0.50', or '0.25' to control the size, and the second controls the input image size, either @@ -114,7 +107,6 @@ import numpy as np from six.moves import urllib import tensorflow as tf -from tensorflow.contrib.quantize.python import quant_ops from tensorflow.python.framework import graph_util from tensorflow.python.framework import tensor_shape from tensorflow.python.platform import gfile @@ -279,7 +271,6 @@ def create_model_graph(model_info): """ with tf.Graph().as_default() as graph: model_path = os.path.join(FLAGS.model_dir, model_info['model_file_name']) - print('Model path: ', model_path) with gfile.FastGFile(model_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) @@ -346,10 +337,7 @@ def maybe_download_and_extract(data_url): statinfo = os.stat(filepath) tf.logging.info('Successfully downloaded', filename, statinfo.st_size, 'bytes.') - print('Extracting file from ', filepath) - tarfile.open(filepath, 'r:gz').extractall(dest_directory) - else: - print('Not extracting or downloading files, model already present in disk') + tarfile.open(filepath, 'r:gz').extractall(dest_directory) def ensure_dir_exists(dir_name): @@ -745,7 +733,7 @@ def variable_summaries(var): def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, - bottleneck_tensor_size, quantize_layer): + bottleneck_tensor_size): """Adds a new softmax and fully-connected layer for training. We need to retrain the top layer to identify our new classes, so this function @@ -757,12 +745,10 @@ def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, Args: class_count: Integer of how many categories of things we're trying to - recognize. + recognize. final_tensor_name: Name string for the new final node that produces results. bottleneck_tensor: The output of the main CNN graph. bottleneck_tensor_size: How many entries in the bottleneck vector. - quantize_layer: Boolean, specifying whether the newly added layer should be - quantized. Returns: The tensors for the training and cross entropy results, and tensors for the @@ -785,41 +771,18 @@ def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, with tf.name_scope('weights'): initial_value = tf.truncated_normal( [bottleneck_tensor_size, class_count], stddev=0.001) + layer_weights = tf.Variable(initial_value, name='final_weights') - if quantize_layer: - quantized_layer_weights = quant_ops.MovingAvgQuantize( - layer_weights, is_training=True) - variable_summaries(quantized_layer_weights) variable_summaries(layer_weights) with tf.name_scope('biases'): layer_biases = tf.Variable(tf.zeros([class_count]), name='final_biases') - if quantize_layer: - quantized_layer_biases = quant_ops.MovingAvgQuantize( - layer_biases, is_training=True) - variable_summaries(quantized_layer_biases) - variable_summaries(layer_biases) - with tf.name_scope('Wx_plus_b'): - if quantize_layer: - logits = tf.matmul(bottleneck_input, - quantized_layer_weights) + quantized_layer_biases - logits = quant_ops.MovingAvgQuantize( - logits, - init_min=-32.0, - init_max=32.0, - is_training=True, - num_bits=8, - narrow_range=False, - ema_decay=0.5) - tf.summary.histogram('pre_activations', logits) - else: - logits = tf.matmul(bottleneck_input, layer_weights) + layer_biases - tf.summary.histogram('pre_activations', logits) + logits = tf.matmul(bottleneck_input, layer_weights) + layer_biases + tf.summary.histogram('pre_activations', logits) final_tensor = tf.nn.softmax(logits, name=final_tensor_name) - tf.summary.histogram('activations', final_tensor) with tf.name_scope('cross_entropy'): @@ -827,7 +790,6 @@ def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor, labels=ground_truth_input, logits=logits) with tf.name_scope('total'): cross_entropy_mean = tf.reduce_mean(cross_entropy) - tf.summary.scalar('cross_entropy', cross_entropy_mean) with tf.name_scope('train'): @@ -863,7 +825,6 @@ def add_evaluation_step(result_tensor, ground_truth_tensor): def save_graph_to_file(sess, graph, graph_file_name): output_graph_def = graph_util.convert_variables_to_constants( sess, graph.as_graph_def(), [FLAGS.final_tensor_name]) - with gfile.FastGFile(graph_file_name, 'wb') as f: f.write(output_graph_def.SerializeToString()) return @@ -897,7 +858,6 @@ def create_model_info(architecture): ValueError: If architecture name is unknown. """ architecture = architecture.lower() - is_quantized = False if architecture == 'inception_v3': # pylint: disable=line-too-long data_url = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz' @@ -942,28 +902,19 @@ def create_model_info(architecture): architecture) return None is_quantized = True - - if is_quantized: - data_url = 'http://download.tensorflow.org/models/mobilenet_v1_' - data_url += version_string + '_' + size_string + '_quantized_frozen.tgz' - bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0' - resized_input_tensor_name = 'Placeholder:0' - model_dir_name = ('mobilenet_v1_' + version_string + '_' + size_string + - '_quantized_frozen') - model_base_name = 'quantized_frozen_graph.pb' - - else: - data_url = 'http://download.tensorflow.org/models/mobilenet_v1_' - data_url += version_string + '_' + size_string + '_frozen.tgz' - bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0' - resized_input_tensor_name = 'input:0' - model_dir_name = 'mobilenet_v1_' + version_string + '_' + size_string - model_base_name = 'frozen_graph.pb' - + data_url = 'http://download.tensorflow.org/models/mobilenet_v1_' + data_url += version_string + '_' + size_string + '_frozen.tgz' + bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0' bottleneck_tensor_size = 1001 input_width = int(size_string) input_height = int(size_string) input_depth = 3 + resized_input_tensor_name = 'input:0' + if is_quantized: + model_base_name = 'quantized_graph.pb' + else: + model_base_name = 'frozen_graph.pb' + model_dir_name = 'mobilenet_v1_' + version_string + '_' + size_string model_file_name = os.path.join(model_dir_name, model_base_name) input_mean = 127.5 input_std = 127.5 @@ -982,7 +933,6 @@ def create_model_info(architecture): 'model_file_name': model_file_name, 'input_mean': input_mean, 'input_std': input_std, - 'quantize_layer': is_quantized, } @@ -1078,7 +1028,7 @@ def main(_): (train_step, cross_entropy, bottleneck_input, ground_truth_input, final_tensor) = add_final_training_ops( len(image_lists.keys()), FLAGS.final_tensor_name, bottleneck_tensor, - model_info['bottleneck_tensor_size'], model_info['quantize_layer']) + model_info['bottleneck_tensor_size']) # Create the operations we need to evaluate the accuracy of our new layer. evaluation_step, prediction = add_evaluation_step( |