· 8 years ago · Feb 02, 2018, 12:32 PM
1# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15r"""Simple transfer learning with Inception v3 or Mobilenet models.
16
17With support for TensorBoard.
18
19This example shows how to take a Inception v3 or Mobilenet model trained on
20ImageNet images, and train a new top layer that can recognize other classes of
21images.
22
23The top layer receives as input a 2048-dimensional vector (1001-dimensional for
24Mobilenet) for each image. We train a softmax layer on top of this
25representation. Assuming the softmax layer contains N labels, this corresponds
26to learning N + 2048*N (or 1001*N) model parameters corresponding to the
27learned biases and weights.
28
29Here's an example, which assumes you have a folder containing class-named
30subfolders, each full of images for each label. The example folder flower_photos
31should have a structure like this:
32
33~/flower_photos/daisy/photo1.jpg
34~/flower_photos/daisy/photo2.jpg
35...
36~/flower_photos/rose/anotherphoto77.jpg
37...
38~/flower_photos/sunflower/somepicture.jpg
39
40The subfolder names are important, since they define what label is applied to
41each image, but the filenames themselves don't matter. Once your images are
42prepared, you can run the training with a command like this:
43
44
45```bash
46bazel build tensorflow/examples/image_retraining:retrain && \
47bazel-bin/tensorflow/examples/image_retraining/retrain \
48 --image_dir ~/flower_photos
49```
50
51Or, if you have a pip installation of tensorflow, `retrain.py` can be run
52without bazel:
53
54```bash
55python tensorflow/examples/image_retraining/retrain.py \
56 --image_dir ~/flower_photos
57```
58
59You can replace the image_dir argument with any folder containing subfolders of
60images. The label for each image is taken from the name of the subfolder it's
61in.
62
63This produces a new model file that can be loaded and run by any TensorFlow
64program, for example the label_image sample code.
65
66By default this script will use the high accuracy, but comparatively large and
67slow Inception v3 model architecture. It's recommended that you start with this
68to validate that you have gathered good training data, but if you want to deploy
69on resource-limited platforms, you can try the `--architecture` flag with a
70Mobilenet model. For example:
71
72Run floating-point version of mobilenet:
73```bash
74python tensorflow/examples/image_retraining/retrain.py \
75 --image_dir ~/flower_photos --architecture mobilenet_1.0_224
76```
77
78Run quantized version of mobilenet:
79```bash
80python tensorflow/examples/image_retraining/retrain.py \
81 --image_dir ~/flower_photos/ --architecture mobilenet_1.0_224_quantized
82```
83
84There are 32 different Mobilenet models to choose from, with a variety of file
85size and latency options. The first number can be '1.0', '0.75', '0.50', or
86'0.25' to control the size, and the second controls the input image size, either
87'224', '192', '160', or '128', with smaller sizes running faster. See
88https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html
89for more information on Mobilenet.
90
91To use with TensorBoard:
92
93By default, this script will log summaries to /tmp/retrain_logs directory
94
95Visualize the summaries with this command:
96
97tensorboard --logdir /tmp/retrain_logs
98
99To use with Tensorflow Serving:
100
101tensorflow_model_server --port=9000 --model_name=inception --model_base_path=/tmp/saved_models/
102
103"""
104from __future__ import absolute_import
105from __future__ import division
106from __future__ import print_function
107
108import argparse
109from datetime import datetime
110import hashlib
111import os.path
112import random
113import re
114import sys
115import tarfile
116
117import numpy as np
118from six.moves import urllib
119import tensorflow as tf
120
121from tensorflow.contrib.quantize.python import quant_ops
122from tensorflow.python.framework import graph_util
123from tensorflow.python.framework import tensor_shape
124from tensorflow.python.platform import gfile
125from tensorflow.python.util import compat
126
127FLAGS = None
128
129# These are all parameters that are tied to the particular model architecture
130# we're using for Inception v3. These include things like tensor names and their
131# sizes. If you want to adapt this script to work with another model, you will
132# need to update these to reflect the values in the network you're using.
133MAX_NUM_IMAGES_PER_CLASS = 2 ** 27 - 1 # ~134M
134
135
136def create_image_lists(image_dir, testing_percentage, validation_percentage):
137 """Builds a list of training images from the file system.
138
139 Analyzes the sub folders in the image directory, splits them into stable
140 training, testing, and validation sets, and returns a data structure
141 describing the lists of images for each label and their paths.
142
143 Args:
144 image_dir: String path to a folder containing subfolders of images.
145 testing_percentage: Integer percentage of the images to reserve for tests.
146 validation_percentage: Integer percentage of images reserved for validation.
147
148 Returns:
149 A dictionary containing an entry for each label subfolder, with images split
150 into training, testing, and validation sets within each label.
151 """
152 if not gfile.Exists(image_dir):
153 tf.logging.error("Image directory '" + image_dir + "' not found.")
154 return None
155 result = {}
156 sub_dirs = [x[0] for x in gfile.Walk(image_dir)]
157 # The root directory comes first, so skip it.
158 is_root_dir = True
159 for sub_dir in sub_dirs:
160 if is_root_dir:
161 is_root_dir = False
162 continue
163 extensions = ['jpg', 'jpeg', 'JPG', 'JPEG']
164 file_list = []
165 dir_name = os.path.basename(sub_dir)
166 if dir_name == image_dir:
167 continue
168 tf.logging.info("Looking for images in '" + dir_name + "'")
169 for extension in extensions:
170 file_glob = os.path.join(image_dir, dir_name, '*.' + extension)
171 file_list.extend(gfile.Glob(file_glob))
172 if not file_list:
173 tf.logging.warning('No files found')
174 continue
175 if len(file_list) < 20:
176 tf.logging.warning(
177 'WARNING: Folder has less than 20 images, which may cause issues.')
178 elif len(file_list) > MAX_NUM_IMAGES_PER_CLASS:
179 tf.logging.warning(
180 'WARNING: Folder {} has more than {} images. Some images will '
181 'never be selected.'.format(dir_name, MAX_NUM_IMAGES_PER_CLASS))
182 label_name = re.sub(r'[^a-z0-9]+', ' ', dir_name.lower())
183 training_images = []
184 testing_images = []
185 validation_images = []
186 for file_name in file_list:
187 base_name = os.path.basename(file_name)
188 # We want to ignore anything after '_nohash_' in the file name when
189 # deciding which set to put an image in, the data set creator has a way of
190 # grouping photos that are close variations of each other. For example
191 # this is used in the plant disease data set to group multiple pictures of
192 # the same leaf.
193 hash_name = re.sub(r'_nohash_.*$', '', file_name)
194 # This looks a bit magical, but we need to decide whether this file should
195 # go into the training, testing, or validation sets, and we want to keep
196 # existing files in the same set even if more files are subsequently
197 # added.
198 # To do that, we need a stable way of deciding based on just the file name
199 # itself, so we do a hash of that and then use that to generate a
200 # probability value that we use to assign it.
201 hash_name_hashed = hashlib.sha1(compat.as_bytes(hash_name)).hexdigest()
202 percentage_hash = ((int(hash_name_hashed, 16) %
203 (MAX_NUM_IMAGES_PER_CLASS + 1)) *
204 (100.0 / MAX_NUM_IMAGES_PER_CLASS))
205 if percentage_hash < validation_percentage:
206 validation_images.append(base_name)
207 elif percentage_hash < (testing_percentage + validation_percentage):
208 testing_images.append(base_name)
209 else:
210 training_images.append(base_name)
211 result[label_name] = {
212 'dir': dir_name,
213 'training': training_images,
214 'testing': testing_images,
215 'validation': validation_images,
216 }
217 return result
218
219
220def get_image_path(image_lists, label_name, index, image_dir, category):
221 """"Returns a path to an image for a label at the given index.
222
223 Args:
224 image_lists: Dictionary of training images for each label.
225 label_name: Label string we want to get an image for.
226 index: Int offset of the image we want. This will be moduloed by the
227 available number of images for the label, so it can be arbitrarily large.
228 image_dir: Root folder string of the subfolders containing the training
229 images.
230 category: Name string of set to pull images from - training, testing, or
231 validation.
232
233 Returns:
234 File system path string to an image that meets the requested parameters.
235
236 """
237 if label_name not in image_lists:
238 tf.logging.fatal('Label does not exist %s.', label_name)
239 label_lists = image_lists[label_name]
240 if category not in label_lists:
241 tf.logging.fatal('Category does not exist %s.', category)
242 category_list = label_lists[category]
243 if not category_list:
244 tf.logging.fatal('Label %s has no images in the category %s.',
245 label_name, category)
246 mod_index = index % len(category_list)
247 base_name = category_list[mod_index]
248 sub_dir = label_lists['dir']
249 full_path = os.path.join(image_dir, sub_dir, base_name)
250 return full_path
251
252
253def get_bottleneck_path(image_lists, label_name, index, bottleneck_dir,
254 category, architecture):
255 """"Returns a path to a bottleneck file for a label at the given index.
256
257 Args:
258 image_lists: Dictionary of training images for each label.
259 label_name: Label string we want to get an image for.
260 index: Integer offset of the image we want. This will be moduloed by the
261 available number of images for the label, so it can be arbitrarily large.
262 bottleneck_dir: Folder string holding cached files of bottleneck values.
263 category: Name string of set to pull images from - training, testing, or
264 validation.
265 architecture: The name of the model architecture.
266
267 Returns:
268 File system path string to an image that meets the requested parameters.
269 """
270 return get_image_path(image_lists, label_name, index, bottleneck_dir,
271 category) + '_' + architecture + '.txt'
272
273
274def create_model_graph(model_info):
275 """"Creates a graph from saved GraphDef file and returns a Graph object.
276
277 Args:
278 model_info: Dictionary containing information about the model architecture.
279
280 Returns:
281 Graph holding the trained Inception network, and various tensors we'll be
282 manipulating.
283 """
284 with tf.Graph().as_default() as graph:
285 model_path = os.path.join(FLAGS.model_dir, model_info['model_file_name'])
286 print('Model path: ', model_path)
287 with gfile.FastGFile(model_path, 'rb') as f:
288 graph_def = tf.GraphDef()
289 graph_def.ParseFromString(f.read())
290 bottleneck_tensor, resized_input_tensor = (tf.import_graph_def(
291 graph_def,
292 name='',
293 return_elements=[
294 model_info['bottleneck_tensor_name'],
295 model_info['resized_input_tensor_name'],
296 ]))
297 return graph, bottleneck_tensor, resized_input_tensor
298
299
300def run_bottleneck_on_image(sess, image_data, image_data_tensor,
301 decoded_image_tensor, resized_input_tensor,
302 bottleneck_tensor):
303 """Runs inference on an image to extract the 'bottleneck' summary layer.
304
305 Args:
306 sess: Current active TensorFlow Session.
307 image_data: String of raw JPEG data.
308 image_data_tensor: Input data layer in the graph.
309 decoded_image_tensor: Output of initial image resizing and preprocessing.
310 resized_input_tensor: The input node of the recognition graph.
311 bottleneck_tensor: Layer before the final softmax.
312
313 Returns:
314 Numpy array of bottleneck values.
315 """
316 # First decode the JPEG image, resize it, and rescale the pixel values.
317 resized_input_values = sess.run(decoded_image_tensor,
318 {image_data_tensor: image_data})
319 # Then run it through the recognition network.
320 bottleneck_values = sess.run(bottleneck_tensor,
321 {resized_input_tensor: resized_input_values})
322 bottleneck_values = np.squeeze(bottleneck_values)
323 return bottleneck_values
324
325
326def maybe_download_and_extract(data_url):
327 """Download and extract model tar file.
328
329 If the pretrained model we're using doesn't already exist, this function
330 downloads it from the TensorFlow.org website and unpacks it into a directory.
331
332 Args:
333 data_url: Web location of the tar file containing the pretrained model.
334 """
335 dest_directory = FLAGS.model_dir
336 if not os.path.exists(dest_directory):
337 os.makedirs(dest_directory)
338 filename = data_url.split('/')[-1]
339 filepath = os.path.join(dest_directory, filename)
340 if not os.path.exists(filepath):
341
342 def _progress(count, block_size, total_size):
343 sys.stdout.write('\r>> Downloading %s %.1f%%' %
344 (filename,
345 float(count * block_size) / float(total_size) * 100.0))
346 sys.stdout.flush()
347
348 filepath, _ = urllib.request.urlretrieve(data_url, filepath, _progress)
349 print()
350 statinfo = os.stat(filepath)
351 tf.logging.info('Successfully downloaded', filename, statinfo.st_size,
352 'bytes.')
353 print('Extracting file from ', filepath)
354 tarfile.open(filepath, 'r:gz').extractall(dest_directory)
355 else:
356 print('Not extracting or downloading files, model already present in disk')
357
358
359def ensure_dir_exists(dir_name):
360 """Makes sure the folder exists on disk.
361
362 Args:
363 dir_name: Path string to the folder we want to create.
364 """
365 if not os.path.exists(dir_name):
366 os.makedirs(dir_name)
367
368
369bottleneck_path_2_bottleneck_values = {}
370
371
372def create_bottleneck_file(bottleneck_path, image_lists, label_name, index,
373 image_dir, category, sess, jpeg_data_tensor,
374 decoded_image_tensor, resized_input_tensor,
375 bottleneck_tensor):
376 """Create a single bottleneck file."""
377 tf.logging.info('Creating bottleneck at ' + bottleneck_path)
378 image_path = get_image_path(image_lists, label_name, index,
379 image_dir, category)
380 if not gfile.Exists(image_path):
381 tf.logging.fatal('File does not exist %s', image_path)
382 image_data = gfile.FastGFile(image_path, 'rb').read()
383 try:
384 bottleneck_values = run_bottleneck_on_image(
385 sess, image_data, jpeg_data_tensor, decoded_image_tensor,
386 resized_input_tensor, bottleneck_tensor)
387 except Exception as e:
388 raise RuntimeError('Error during processing file %s (%s)' % (image_path,
389 str(e)))
390 bottleneck_string = ','.join(str(x) for x in bottleneck_values)
391 with open(bottleneck_path, 'w') as bottleneck_file:
392 bottleneck_file.write(bottleneck_string)
393
394
395def get_or_create_bottleneck(sess, image_lists, label_name, index, image_dir,
396 category, bottleneck_dir, jpeg_data_tensor,
397 decoded_image_tensor, resized_input_tensor,
398 bottleneck_tensor, architecture):
399 """Retrieves or calculates bottleneck values for an image.
400
401 If a cached version of the bottleneck data exists on-disk, return that,
402 otherwise calculate the data and save it to disk for future use.
403
404 Args:
405 sess: The current active TensorFlow Session.
406 image_lists: Dictionary of training images for each label.
407 label_name: Label string we want to get an image for.
408 index: Integer offset of the image we want. This will be modulo-ed by the
409 available number of images for the label, so it can be arbitrarily large.
410 image_dir: Root folder string of the subfolders containing the training
411 images.
412 category: Name string of which set to pull images from - training, testing,
413 or validation.
414 bottleneck_dir: Folder string holding cached files of bottleneck values.
415 jpeg_data_tensor: The tensor to feed loaded jpeg data into.
416 decoded_image_tensor: The output of decoding and resizing the image.
417 resized_input_tensor: The input node of the recognition graph.
418 bottleneck_tensor: The output tensor for the bottleneck values.
419 architecture: The name of the model architecture.
420
421 Returns:
422 Numpy array of values produced by the bottleneck layer for the image.
423 """
424 label_lists = image_lists[label_name]
425 sub_dir = label_lists['dir']
426 sub_dir_path = os.path.join(bottleneck_dir, sub_dir)
427 ensure_dir_exists(sub_dir_path)
428 bottleneck_path = get_bottleneck_path(image_lists, label_name, index,
429 bottleneck_dir, category, architecture)
430 if not os.path.exists(bottleneck_path):
431 create_bottleneck_file(bottleneck_path, image_lists, label_name, index,
432 image_dir, category, sess, jpeg_data_tensor,
433 decoded_image_tensor, resized_input_tensor,
434 bottleneck_tensor)
435 with open(bottleneck_path, 'r') as bottleneck_file:
436 bottleneck_string = bottleneck_file.read()
437 did_hit_error = False
438 try:
439 bottleneck_values = [float(x) for x in bottleneck_string.split(',')]
440 except ValueError:
441 tf.logging.warning('Invalid float found, recreating bottleneck')
442 did_hit_error = True
443 if did_hit_error:
444 create_bottleneck_file(bottleneck_path, image_lists, label_name, index,
445 image_dir, category, sess, jpeg_data_tensor,
446 decoded_image_tensor, resized_input_tensor,
447 bottleneck_tensor)
448 with open(bottleneck_path, 'r') as bottleneck_file:
449 bottleneck_string = bottleneck_file.read()
450 # Allow exceptions to propagate here, since they shouldn't happen after a
451 # fresh creation
452 bottleneck_values = [float(x) for x in bottleneck_string.split(',')]
453 return bottleneck_values
454
455
456def cache_bottlenecks(sess, image_lists, image_dir, bottleneck_dir,
457 jpeg_data_tensor, decoded_image_tensor,
458 resized_input_tensor, bottleneck_tensor, architecture):
459 """Ensures all the training, testing, and validation bottlenecks are cached.
460
461 Because we're likely to read the same image multiple times (if there are no
462 distortions applied during training) it can speed things up a lot if we
463 calculate the bottleneck layer values once for each image during
464 preprocessing, and then just read those cached values repeatedly during
465 training. Here we go through all the images we've found, calculate those
466 values, and save them off.
467
468 Args:
469 sess: The current active TensorFlow Session.
470 image_lists: Dictionary of training images for each label.
471 image_dir: Root folder string of the subfolders containing the training
472 images.
473 bottleneck_dir: Folder string holding cached files of bottleneck values.
474 jpeg_data_tensor: Input tensor for jpeg data from file.
475 decoded_image_tensor: The output of decoding and resizing the image.
476 resized_input_tensor: The input node of the recognition graph.
477 bottleneck_tensor: The penultimate output layer of the graph.
478 architecture: The name of the model architecture.
479
480 Returns:
481 Nothing.
482 """
483 how_many_bottlenecks = 0
484 ensure_dir_exists(bottleneck_dir)
485 for label_name, label_lists in image_lists.items():
486 for category in ['training', 'testing', 'validation']:
487 category_list = label_lists[category]
488 for index, unused_base_name in enumerate(category_list):
489 get_or_create_bottleneck(
490 sess, image_lists, label_name, index, image_dir, category,
491 bottleneck_dir, jpeg_data_tensor, decoded_image_tensor,
492 resized_input_tensor, bottleneck_tensor, architecture)
493
494 how_many_bottlenecks += 1
495 if how_many_bottlenecks % 100 == 0:
496 tf.logging.info(
497 str(how_many_bottlenecks) + ' bottleneck files created.')
498
499
500def get_random_cached_bottlenecks(sess, image_lists, how_many, category,
501 bottleneck_dir, image_dir, jpeg_data_tensor,
502 decoded_image_tensor, resized_input_tensor,
503 bottleneck_tensor, architecture):
504 """Retrieves bottleneck values for cached images.
505
506 If no distortions are being applied, this function can retrieve the cached
507 bottleneck values directly from disk for images. It picks a random set of
508 images from the specified category.
509
510 Args:
511 sess: Current TensorFlow Session.
512 image_lists: Dictionary of training images for each label.
513 how_many: If positive, a random sample of this size will be chosen.
514 If negative, all bottlenecks will be retrieved.
515 category: Name string of which set to pull from - training, testing, or
516 validation.
517 bottleneck_dir: Folder string holding cached files of bottleneck values.
518 image_dir: Root folder string of the subfolders containing the training
519 images.
520 jpeg_data_tensor: The layer to feed jpeg image data into.
521 decoded_image_tensor: The output of decoding and resizing the image.
522 resized_input_tensor: The input node of the recognition graph.
523 bottleneck_tensor: The bottleneck output layer of the CNN graph.
524 architecture: The name of the model architecture.
525
526 Returns:
527 List of bottleneck arrays, their corresponding ground truths, and the
528 relevant filenames.
529 """
530 class_count = len(image_lists.keys())
531 bottlenecks = []
532 ground_truths = []
533 filenames = []
534 if how_many >= 0:
535 # Retrieve a random sample of bottlenecks.
536 for unused_i in range(how_many):
537 label_index = random.randrange(class_count)
538 label_name = list(image_lists.keys())[label_index]
539 image_index = random.randrange(MAX_NUM_IMAGES_PER_CLASS + 1)
540 image_name = get_image_path(image_lists, label_name, image_index,
541 image_dir, category)
542 bottleneck = get_or_create_bottleneck(
543 sess, image_lists, label_name, image_index, image_dir, category,
544 bottleneck_dir, jpeg_data_tensor, decoded_image_tensor,
545 resized_input_tensor, bottleneck_tensor, architecture)
546 bottlenecks.append(bottleneck)
547 ground_truths.append(label_index)
548 filenames.append(image_name)
549 else:
550 # Retrieve all bottlenecks.
551 for label_index, label_name in enumerate(image_lists.keys()):
552 for image_index, image_name in enumerate(
553 image_lists[label_name][category]):
554 image_name = get_image_path(image_lists, label_name, image_index,
555 image_dir, category)
556 bottleneck = get_or_create_bottleneck(
557 sess, image_lists, label_name, image_index, image_dir, category,
558 bottleneck_dir, jpeg_data_tensor, decoded_image_tensor,
559 resized_input_tensor, bottleneck_tensor, architecture)
560 bottlenecks.append(bottleneck)
561 ground_truths.append(label_index)
562 filenames.append(image_name)
563 return bottlenecks, ground_truths, filenames
564
565
566def get_random_distorted_bottlenecks(
567 sess, image_lists, how_many, category, image_dir, input_jpeg_tensor,
568 distorted_image, resized_input_tensor, bottleneck_tensor):
569 """Retrieves bottleneck values for training images, after distortions.
570
571 If we're training with distortions like crops, scales, or flips, we have to
572 recalculate the full model for every image, and so we can't use cached
573 bottleneck values. Instead we find random images for the requested category,
574 run them through the distortion graph, and then the full graph to get the
575 bottleneck results for each.
576
577 Args:
578 sess: Current TensorFlow Session.
579 image_lists: Dictionary of training images for each label.
580 how_many: The integer number of bottleneck values to return.
581 category: Name string of which set of images to fetch - training, testing,
582 or validation.
583 image_dir: Root folder string of the subfolders containing the training
584 images.
585 input_jpeg_tensor: The input layer we feed the image data to.
586 distorted_image: The output node of the distortion graph.
587 resized_input_tensor: The input node of the recognition graph.
588 bottleneck_tensor: The bottleneck output layer of the CNN graph.
589
590 Returns:
591 List of bottleneck arrays and their corresponding ground truths.
592 """
593 class_count = len(image_lists.keys())
594 bottlenecks = []
595 ground_truths = []
596 for unused_i in range(how_many):
597 label_index = random.randrange(class_count)
598 label_name = list(image_lists.keys())[label_index]
599 image_index = random.randrange(MAX_NUM_IMAGES_PER_CLASS + 1)
600 image_path = get_image_path(image_lists, label_name, image_index, image_dir,
601 category)
602 if not gfile.Exists(image_path):
603 tf.logging.fatal('File does not exist %s', image_path)
604 jpeg_data = gfile.FastGFile(image_path, 'rb').read()
605 # Note that we materialize the distorted_image_data as a numpy array before
606 # sending running inference on the image. This involves 2 memory copies and
607 # might be optimized in other implementations.
608 distorted_image_data = sess.run(distorted_image,
609 {input_jpeg_tensor: jpeg_data})
610 bottleneck_values = sess.run(bottleneck_tensor,
611 {resized_input_tensor: distorted_image_data})
612 bottleneck_values = np.squeeze(bottleneck_values)
613 bottlenecks.append(bottleneck_values)
614 ground_truths.append(label_index)
615 return bottlenecks, ground_truths
616
617
618def should_distort_images(flip_left_right, random_crop, random_scale,
619 random_brightness):
620 """Whether any distortions are enabled, from the input flags.
621
622 Args:
623 flip_left_right: Boolean whether to randomly mirror images horizontally.
624 random_crop: Integer percentage setting the total margin used around the
625 crop box.
626 random_scale: Integer percentage of how much to vary the scale by.
627 random_brightness: Integer range to randomly multiply the pixel values by.
628
629 Returns:
630 Boolean value indicating whether any distortions should be applied.
631 """
632 return (flip_left_right or (random_crop != 0) or (random_scale != 0) or
633 (random_brightness != 0))
634
635
636def add_input_distortions(flip_left_right, random_crop, random_scale,
637 random_brightness, input_width, input_height,
638 input_depth, input_mean, input_std):
639 """Creates the operations to apply the specified distortions.
640
641 During training it can help to improve the results if we run the images
642 through simple distortions like crops, scales, and flips. These reflect the
643 kind of variations we expect in the real world, and so can help train the
644 model to cope with natural data more effectively. Here we take the supplied
645 parameters and construct a network of operations to apply them to an image.
646
647 Cropping
648 ~~~~~~~~
649
650 Cropping is done by placing a bounding box at a random position in the full
651 image. The cropping parameter controls the size of that box relative to the
652 input image. If it's zero, then the box is the same size as the input and no
653 cropping is performed. If the value is 50%, then the crop box will be half the
654 width and height of the input. In a diagram it looks like this:
655
656 < width >
657 +---------------------+
658 | |
659 | width - crop% |
660 | < > |
661 | +------+ |
662 | | | |
663 | | | |
664 | | | |
665 | +------+ |
666 | |
667 | |
668 +---------------------+
669
670 Scaling
671 ~~~~~~~
672
673 Scaling is a lot like cropping, except that the bounding box is always
674 centered and its size varies randomly within the given range. For example if
675 the scale percentage is zero, then the bounding box is the same size as the
676 input and no scaling is applied. If it's 50%, then the bounding box will be in
677 a random range between half the width and height and full size.
678
679 Args:
680 flip_left_right: Boolean whether to randomly mirror images horizontally.
681 random_crop: Integer percentage setting the total margin used around the
682 crop box.
683 random_scale: Integer percentage of how much to vary the scale by.
684 random_brightness: Integer range to randomly multiply the pixel values by.
685 graph.
686 input_width: Horizontal size of expected input image to model.
687 input_height: Vertical size of expected input image to model.
688 input_depth: How many channels the expected input image should have.
689 input_mean: Pixel value that should be zero in the image for the graph.
690 input_std: How much to divide the pixel values by before recognition.
691
692 Returns:
693 The jpeg input layer and the distorted result tensor.
694 """
695
696 jpeg_data = tf.placeholder(tf.string, name='DistortJPGInput')
697 decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
698 decoded_image_as_float = tf.cast(decoded_image, dtype=tf.float32)
699 decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
700 margin_scale = 1.0 + (random_crop / 100.0)
701 resize_scale = 1.0 + (random_scale / 100.0)
702 margin_scale_value = tf.constant(margin_scale)
703 resize_scale_value = tf.random_uniform(tensor_shape.scalar(),
704 minval=1.0,
705 maxval=resize_scale)
706 scale_value = tf.multiply(margin_scale_value, resize_scale_value)
707 precrop_width = tf.multiply(scale_value, input_width)
708 precrop_height = tf.multiply(scale_value, input_height)
709 precrop_shape = tf.stack([precrop_height, precrop_width])
710 precrop_shape_as_int = tf.cast(precrop_shape, dtype=tf.int32)
711 precropped_image = tf.image.resize_bilinear(decoded_image_4d,
712 precrop_shape_as_int)
713 precropped_image_3d = tf.squeeze(precropped_image, squeeze_dims=[0])
714 cropped_image = tf.random_crop(precropped_image_3d,
715 [input_height, input_width, input_depth])
716 if flip_left_right:
717 flipped_image = tf.image.random_flip_left_right(cropped_image)
718 else:
719 flipped_image = cropped_image
720 brightness_min = 1.0 - (random_brightness / 100.0)
721 brightness_max = 1.0 + (random_brightness / 100.0)
722 brightness_value = tf.random_uniform(tensor_shape.scalar(),
723 minval=brightness_min,
724 maxval=brightness_max)
725 brightened_image = tf.multiply(flipped_image, brightness_value)
726 offset_image = tf.subtract(brightened_image, input_mean)
727 mul_image = tf.multiply(offset_image, 1.0 / input_std)
728 distort_result = tf.expand_dims(mul_image, 0, name='DistortResult')
729 return jpeg_data, distort_result
730
731
732def variable_summaries(var):
733 """Attach a lot of summaries to a Tensor (for TensorBoard visualization)."""
734 with tf.name_scope('summaries'):
735 mean = tf.reduce_mean(var)
736 tf.summary.scalar('mean', mean)
737 with tf.name_scope('stddev'):
738 stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
739 tf.summary.scalar('stddev', stddev)
740 tf.summary.scalar('max', tf.reduce_max(var))
741 tf.summary.scalar('min', tf.reduce_min(var))
742 tf.summary.histogram('histogram', var)
743
744
745def add_final_training_ops(class_count, final_tensor_name, bottleneck_tensor,
746 bottleneck_tensor_size, quantize_layer):
747 """Adds a new softmax and fully-connected layer for training.
748
749 We need to retrain the top layer to identify our new classes, so this function
750 adds the right operations to the graph, along with some variables to hold the
751 weights, and then sets up all the gradients for the backward pass.
752
753 The set up for the softmax and fully-connected layers is based on:
754 https://www.tensorflow.org/versions/master/tutorials/mnist/beginners/index.html
755
756 Args:
757 class_count: Integer of how many categories of things we're trying to
758 recognize.
759 final_tensor_name: Name string for the new final node that produces results.
760 bottleneck_tensor: The output of the main CNN graph.
761 bottleneck_tensor_size: How many entries in the bottleneck vector.
762 quantize_layer: Boolean, specifying whether the newly added layer should be
763 quantized.
764
765 Returns:
766 The tensors for the training and cross entropy results, and tensors for the
767 bottleneck input and ground truth input.
768 """
769 with tf.name_scope('input'):
770 bottleneck_input = tf.placeholder_with_default(
771 bottleneck_tensor,
772 shape=[None, bottleneck_tensor_size],
773 name='BottleneckInputPlaceholder')
774
775 ground_truth_input = tf.placeholder(
776 tf.int64, [None], name='GroundTruthInput')
777
778 # Organizing the following ops as `final_training_ops` so they're easier
779 # to see in TensorBoard
780 layer_name = 'final_training_ops'
781 with tf.name_scope(layer_name):
782 with tf.name_scope('weights'):
783 initial_value = tf.truncated_normal(
784 [bottleneck_tensor_size, class_count], stddev=0.001)
785 layer_weights = tf.Variable(initial_value, name='final_weights')
786 if quantize_layer:
787 quantized_layer_weights = quant_ops.MovingAvgQuantize(
788 layer_weights, is_training=True)
789 variable_summaries(quantized_layer_weights)
790
791 variable_summaries(layer_weights)
792 with tf.name_scope('biases'):
793 layer_biases = tf.Variable(tf.zeros([class_count]), name='final_biases')
794 if quantize_layer:
795 quantized_layer_biases = quant_ops.MovingAvgQuantize(
796 layer_biases, is_training=True)
797 variable_summaries(quantized_layer_biases)
798
799 variable_summaries(layer_biases)
800
801 with tf.name_scope('Wx_plus_b'):
802 if quantize_layer:
803 logits = tf.matmul(bottleneck_input,
804 quantized_layer_weights) + quantized_layer_biases
805 logits = quant_ops.MovingAvgQuantize(
806 logits,
807 init_min=-32.0,
808 init_max=32.0,
809 is_training=True,
810 num_bits=8,
811 narrow_range=False,
812 ema_decay=0.5)
813 tf.summary.histogram('pre_activations', logits)
814 else:
815 logits = tf.matmul(bottleneck_input, layer_weights) + layer_biases
816 tf.summary.histogram('pre_activations', logits)
817 scores = tf.identity(logits, name="scores")
818 final_tensor = tf.nn.softmax(logits, name=final_tensor_name)
819
820 tf.summary.histogram('activations', final_tensor)
821
822 with tf.name_scope('cross_entropy'):
823 cross_entropy_mean = tf.losses.sparse_softmax_cross_entropy(
824 labels=ground_truth_input, logits=logits)
825
826 tf.summary.scalar('cross_entropy', cross_entropy_mean)
827
828 with tf.name_scope('train'):
829 optimizer = tf.train.GradientDescentOptimizer(FLAGS.learning_rate)
830 train_step = optimizer.minimize(cross_entropy_mean)
831
832 return (train_step, cross_entropy_mean, bottleneck_input, ground_truth_input,
833 final_tensor)
834
835
836def add_evaluation_step(result_tensor, ground_truth_tensor):
837 """Inserts the operations we need to evaluate the accuracy of our results.
838
839 Args:
840 result_tensor: The new final node that produces results.
841 ground_truth_tensor: The node we feed ground truth data
842 into.
843
844 Returns:
845 Tuple of (evaluation step, prediction).
846 """
847 with tf.name_scope('accuracy'):
848 with tf.name_scope('correct_prediction'):
849 prediction = tf.argmax(result_tensor, 1)
850 correct_prediction = tf.equal(prediction, ground_truth_tensor)
851 with tf.name_scope('accuracy'):
852 evaluation_step = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
853 tf.summary.scalar('accuracy', evaluation_step)
854 return evaluation_step, prediction
855
856
857def save_graph_to_file(sess, graph, graph_file_name):
858 output_graph_def = graph_util.convert_variables_to_constants(
859 sess, graph.as_graph_def(), [FLAGS.final_tensor_name])
860
861 with gfile.FastGFile(graph_file_name, 'wb') as f:
862 f.write(output_graph_def.SerializeToString())
863 return
864
865
866def prepare_file_system():
867 # Setup the directory we'll write summaries to for TensorBoard
868 if tf.gfile.Exists(FLAGS.summaries_dir):
869 tf.gfile.DeleteRecursively(FLAGS.summaries_dir)
870 tf.gfile.MakeDirs(FLAGS.summaries_dir)
871 if FLAGS.intermediate_store_frequency > 0:
872 ensure_dir_exists(FLAGS.intermediate_output_graphs_dir)
873 return
874
875
876def create_model_info(architecture):
877 """Given the name of a model architecture, returns information about it.
878
879 There are different base image recognition pretrained models that can be
880 retrained using transfer learning, and this function translates from the name
881 of a model to the attributes that are needed to download and train with it.
882
883 Args:
884 architecture: Name of a model architecture.
885
886 Returns:
887 Dictionary of information about the model, or None if the name isn't
888 recognized
889
890 Raises:
891 ValueError: If architecture name is unknown.
892 """
893 architecture = architecture.lower()
894 is_quantized = False
895 if architecture == 'inception_v3':
896 # pylint: disable=line-too-long
897 data_url = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'
898 # pylint: enable=line-too-long
899 bottleneck_tensor_name = 'pool_3/_reshape:0'
900 bottleneck_tensor_size = 2048
901 input_width = 299
902 input_height = 299
903 input_depth = 3
904 resized_input_tensor_name = 'Mul:0'
905 model_file_name = 'classify_image_graph_def.pb'
906 input_mean = 128
907 input_std = 128
908 elif architecture.startswith('mobilenet_'):
909 parts = architecture.split('_')
910 if len(parts) != 3 and len(parts) != 4:
911 tf.logging.error("Couldn't understand architecture name '%s'",
912 architecture)
913 return None
914 version_string = parts[1]
915 if (version_string != '1.0' and version_string != '0.75' and
916 version_string != '0.50' and version_string != '0.25'):
917 tf.logging.error(
918 """"The Mobilenet version should be '1.0', '0.75', '0.50', or '0.25',
919 but found '%s' for architecture '%s'""",
920 version_string, architecture)
921 return None
922 size_string = parts[2]
923 if (size_string != '224' and size_string != '192' and
924 size_string != '160' and size_string != '128'):
925 tf.logging.error(
926 """The Mobilenet input size should be '224', '192', '160', or '128',
927 but found '%s' for architecture '%s'""",
928 size_string, architecture)
929 return None
930 if len(parts) == 3:
931 is_quantized = False
932 else:
933 if parts[3] != 'quantized':
934 tf.logging.error(
935 "Couldn't understand architecture suffix '%s' for '%s'", parts[3],
936 architecture)
937 return None
938 is_quantized = True
939
940 if is_quantized:
941 data_url = 'http://download.tensorflow.org/models/mobilenet_v1_'
942 data_url += version_string + '_' + size_string + '_quantized_frozen.tgz'
943 bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0'
944 resized_input_tensor_name = 'Placeholder:0'
945 model_dir_name = ('mobilenet_v1_' + version_string + '_' + size_string +
946 '_quantized_frozen')
947 model_base_name = 'quantized_frozen_graph.pb'
948
949 else:
950 data_url = 'http://download.tensorflow.org/models/mobilenet_v1_'
951 data_url += version_string + '_' + size_string + '_frozen.tgz'
952 bottleneck_tensor_name = 'MobilenetV1/Predictions/Reshape:0'
953 resized_input_tensor_name = 'input:0'
954 model_dir_name = 'mobilenet_v1_' + version_string + '_' + size_string
955 model_base_name = 'frozen_graph.pb'
956
957 bottleneck_tensor_size = 1001
958 input_width = int(size_string)
959 input_height = int(size_string)
960 input_depth = 3
961 model_file_name = os.path.join(model_dir_name, model_base_name)
962 input_mean = 127.5
963 input_std = 127.5
964 else:
965 tf.logging.error("Couldn't understand architecture name '%s'", architecture)
966 raise ValueError('Unknown architecture', architecture)
967
968 return {
969 'data_url': data_url,
970 'bottleneck_tensor_name': bottleneck_tensor_name,
971 'bottleneck_tensor_size': bottleneck_tensor_size,
972 'input_width': input_width,
973 'input_height': input_height,
974 'input_depth': input_depth,
975 'resized_input_tensor_name': resized_input_tensor_name,
976 'model_file_name': model_file_name,
977 'input_mean': input_mean,
978 'input_std': input_std,
979 'quantize_layer': is_quantized,
980 }
981
982
983def add_jpeg_decoding(input_width, input_height, input_depth, input_mean,
984 input_std):
985 """Adds operations that perform JPEG decoding and resizing to the graph..
986
987 Args:
988 input_width: Desired width of the image fed into the recognizer graph.
989 input_height: Desired width of the image fed into the recognizer graph.
990 input_depth: Desired channels of the image fed into the recognizer graph.
991 input_mean: Pixel value that should be zero in the image for the graph.
992 input_std: How much to divide the pixel values by before recognition.
993
994 Returns:
995 Tensors for the node to feed JPEG data into, and the output of the
996 preprocessing steps.
997 """
998 jpeg_data = tf.placeholder(tf.string, name='DecodeJPGInput')
999 decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth)
1000 decoded_image_as_float = tf.cast(decoded_image, dtype=tf.float32)
1001 decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0)
1002 resize_shape = tf.stack([input_height, input_width])
1003 resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32)
1004 resized_image = tf.image.resize_bilinear(decoded_image_4d,
1005 resize_shape_as_int)
1006 offset_image = tf.subtract(resized_image, input_mean)
1007 mul_image = tf.multiply(offset_image, 1.0 / input_std)
1008 return jpeg_data, mul_image
1009
1010
1011def export_model(sess, image_lists, architecture, saved_model_dir):
1012 """Exports model for serving.
1013
1014 Args:
1015 sess: Current active TensorFlow Session.
1016 architecture: Model architecture.
1017 saved_model_dir: Directory in which to save exported model and variables.
1018 """
1019 if architecture == 'inception_v3':
1020 input_tensor = 'DecodeJpeg/contents:0'
1021 elif architecture.startswith('mobilenet_'):
1022 input_tensor = 'input:0'
1023 else:
1024 raise ValueError('Unknown architecture', architecture)
1025 in_image = sess.graph.get_tensor_by_name(input_tensor)
1026 inputs = {'images': tf.saved_model.utils.build_tensor_info(in_image)}
1027
1028 #out_classes = sess.graph.get_tensor_by_name('final_result:0')
1029 out_scores = sess.graph.get_tensor_by_name('scores:0')
1030
1031 values, indices = tf.nn.top_k(out_scores, 5)
1032
1033 class_descriptions = []
1034 for s in image_lists.keys():
1035 class_descriptions.append(s)
1036 class_tensor = tf.constant(class_descriptions)
1037
1038 table = tf.contrib.lookup.index_to_string_table_from_tensor(class_tensor)
1039 classes = table.lookup(tf.to_int64(indices))
1040
1041 outputs = {'classes': tf.saved_model.utils.build_tensor_info(classes),
1042 'scores': tf.saved_model.utils.build_tensor_info(values)}
1043
1044 signature = tf.saved_model.signature_def_utils.build_signature_def(
1045 inputs=inputs,
1046 outputs=outputs,
1047 method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME
1048 )
1049
1050 legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
1051
1052 # Save out the SavedModel.
1053 builder = tf.saved_model.builder.SavedModelBuilder(saved_model_dir)
1054 builder.add_meta_graph_and_variables(
1055 sess, [tf.saved_model.tag_constants.SERVING],
1056 signature_def_map={
1057 'predict_images': signature
1058 },
1059 legacy_init_op=legacy_init_op)
1060 builder.save()
1061
1062
1063def main(_):
1064 # Needed to make sure the logging output is visible.
1065 # See https://github.com/tensorflow/tensorflow/issues/3047
1066 tf.logging.set_verbosity(tf.logging.INFO)
1067
1068 # Prepare necessary directories that can be used during training
1069 prepare_file_system()
1070
1071 # Gather information about the model architecture we'll be using.
1072 model_info = create_model_info(FLAGS.architecture)
1073 if not model_info:
1074 tf.logging.error('Did not recognize architecture flag')
1075 return -1
1076
1077 # Set up the pre-trained graph.
1078 maybe_download_and_extract(model_info['data_url'])
1079 graph, bottleneck_tensor, resized_image_tensor = (
1080 create_model_graph(model_info))
1081
1082 # Look at the folder structure, and create lists of all the images.
1083 image_lists = create_image_lists(FLAGS.image_dir, FLAGS.testing_percentage,
1084 FLAGS.validation_percentage)
1085 class_count = len(image_lists.keys())
1086 if class_count == 0:
1087 tf.logging.error('No valid folders of images found at ' + FLAGS.image_dir)
1088 return -1
1089 if class_count == 1:
1090 tf.logging.error('Only one valid folder of images found at ' +
1091 FLAGS.image_dir +
1092 ' - multiple classes are needed for classification.')
1093 return -1
1094
1095 # See if the command-line flags mean we're applying any distortions.
1096 do_distort_images = should_distort_images(
1097 FLAGS.flip_left_right, FLAGS.random_crop, FLAGS.random_scale,
1098 FLAGS.random_brightness)
1099
1100 with tf.Session(graph=graph) as sess:
1101 # Set up the image decoding sub-graph.
1102 jpeg_data_tensor, decoded_image_tensor = add_jpeg_decoding(
1103 model_info['input_width'], model_info['input_height'],
1104 model_info['input_depth'], model_info['input_mean'],
1105 model_info['input_std'])
1106
1107 if do_distort_images:
1108 # We will be applying distortions, so setup the operations we'll need.
1109 (distorted_jpeg_data_tensor,
1110 distorted_image_tensor) = add_input_distortions(
1111 FLAGS.flip_left_right, FLAGS.random_crop, FLAGS.random_scale,
1112 FLAGS.random_brightness, model_info['input_width'],
1113 model_info['input_height'], model_info['input_depth'],
1114 model_info['input_mean'], model_info['input_std'])
1115 else:
1116 # We'll make sure we've calculated the 'bottleneck' image summaries and
1117 # cached them on disk.
1118 cache_bottlenecks(sess, image_lists, FLAGS.image_dir,
1119 FLAGS.bottleneck_dir, jpeg_data_tensor,
1120 decoded_image_tensor, resized_image_tensor,
1121 bottleneck_tensor, FLAGS.architecture)
1122
1123 # Add the new layer that we'll be training.
1124 (train_step, cross_entropy, bottleneck_input, ground_truth_input,
1125 final_tensor) = add_final_training_ops(
1126 len(image_lists.keys()), FLAGS.final_tensor_name, bottleneck_tensor,
1127 model_info['bottleneck_tensor_size'], model_info['quantize_layer'])
1128
1129 # Create the operations we need to evaluate the accuracy of our new layer.
1130 evaluation_step, prediction = add_evaluation_step(
1131 final_tensor, ground_truth_input)
1132
1133 # Merge all the summaries and write them out to the summaries_dir
1134 merged = tf.summary.merge_all()
1135 train_writer = tf.summary.FileWriter(FLAGS.summaries_dir + '/train',
1136 sess.graph)
1137
1138 validation_writer = tf.summary.FileWriter(
1139 FLAGS.summaries_dir + '/validation')
1140
1141 # Set up all our weights to their initial default values.
1142 init = tf.global_variables_initializer()
1143 sess.run(init)
1144
1145 # Run the training for as many cycles as requested on the command line.
1146 for i in range(FLAGS.how_many_training_steps):
1147 # Get a batch of input bottleneck values, either calculated fresh every
1148 # time with distortions applied, or from the cache stored on disk.
1149 if do_distort_images:
1150 (train_bottlenecks,
1151 train_ground_truth) = get_random_distorted_bottlenecks(
1152 sess, image_lists, FLAGS.train_batch_size, 'training',
1153 FLAGS.image_dir, distorted_jpeg_data_tensor,
1154 distorted_image_tensor, resized_image_tensor, bottleneck_tensor)
1155 else:
1156 (train_bottlenecks,
1157 train_ground_truth, _) = get_random_cached_bottlenecks(
1158 sess, image_lists, FLAGS.train_batch_size, 'training',
1159 FLAGS.bottleneck_dir, FLAGS.image_dir, jpeg_data_tensor,
1160 decoded_image_tensor, resized_image_tensor, bottleneck_tensor,
1161 FLAGS.architecture)
1162 # Feed the bottlenecks and ground truth into the graph, and run a training
1163 # step. Capture training summaries for TensorBoard with the `merged` op.
1164 train_summary, _ = sess.run(
1165 [merged, train_step],
1166 feed_dict={bottleneck_input: train_bottlenecks,
1167 ground_truth_input: train_ground_truth})
1168 train_writer.add_summary(train_summary, i)
1169
1170 # Every so often, print out how well the graph is training.
1171 is_last_step = (i + 1 == FLAGS.how_many_training_steps)
1172 if (i % FLAGS.eval_step_interval) == 0 or is_last_step:
1173 train_accuracy, cross_entropy_value = sess.run(
1174 [evaluation_step, cross_entropy],
1175 feed_dict={bottleneck_input: train_bottlenecks,
1176 ground_truth_input: train_ground_truth})
1177 tf.logging.info('%s: Step %d: Train accuracy = %.1f%%' %
1178 (datetime.now(), i, train_accuracy * 100))
1179 tf.logging.info('%s: Step %d: Cross entropy = %f' %
1180 (datetime.now(), i, cross_entropy_value))
1181 validation_bottlenecks, validation_ground_truth, _ = (
1182 get_random_cached_bottlenecks(
1183 sess, image_lists, FLAGS.validation_batch_size, 'validation',
1184 FLAGS.bottleneck_dir, FLAGS.image_dir, jpeg_data_tensor,
1185 decoded_image_tensor, resized_image_tensor, bottleneck_tensor,
1186 FLAGS.architecture))
1187 # Run a validation step and capture training summaries for TensorBoard
1188 # with the `merged` op.
1189 validation_summary, validation_accuracy = sess.run(
1190 [merged, evaluation_step],
1191 feed_dict={bottleneck_input: validation_bottlenecks,
1192 ground_truth_input: validation_ground_truth})
1193 validation_writer.add_summary(validation_summary, i)
1194 tf.logging.info('%s: Step %d: Validation accuracy = %.1f%% (N=%d)' %
1195 (datetime.now(), i, validation_accuracy * 100,
1196 len(validation_bottlenecks)))
1197
1198 # Store intermediate results
1199 intermediate_frequency = FLAGS.intermediate_store_frequency
1200
1201 if (intermediate_frequency > 0 and (i % intermediate_frequency == 0)
1202 and i > 0):
1203 intermediate_file_name = (FLAGS.intermediate_output_graphs_dir +
1204 'intermediate_' + str(i) + '.pb')
1205 tf.logging.info('Save intermediate result to : ' +
1206 intermediate_file_name)
1207 save_graph_to_file(sess, graph, intermediate_file_name)
1208
1209 # We've completed all our training, so run a final test evaluation on
1210 # some new images we haven't used before.
1211 test_bottlenecks, test_ground_truth, test_filenames = (
1212 get_random_cached_bottlenecks(
1213 sess, image_lists, FLAGS.test_batch_size, 'testing',
1214 FLAGS.bottleneck_dir, FLAGS.image_dir, jpeg_data_tensor,
1215 decoded_image_tensor, resized_image_tensor, bottleneck_tensor,
1216 FLAGS.architecture))
1217 test_accuracy, predictions = sess.run(
1218 [evaluation_step, prediction],
1219 feed_dict={bottleneck_input: test_bottlenecks,
1220 ground_truth_input: test_ground_truth})
1221 tf.logging.info('Final test accuracy = %.1f%% (N=%d)' %
1222 (test_accuracy * 100, len(test_bottlenecks)))
1223
1224 if FLAGS.print_misclassified_test_images:
1225 tf.logging.info('=== MISCLASSIFIED TEST IMAGES ===')
1226 for i, test_filename in enumerate(test_filenames):
1227 if predictions[i] != test_ground_truth[i]:
1228 tf.logging.info('%70s %s' %
1229 (test_filename,
1230 list(image_lists.keys())[predictions[i]]))
1231
1232 # Write out the trained graph and labels with the weights stored as
1233 # constants.
1234 save_graph_to_file(sess, graph, FLAGS.output_graph)
1235 with gfile.FastGFile(FLAGS.output_labels, 'w') as f:
1236 f.write('\n'.join(image_lists.keys()) + '\n')
1237
1238 export_model(sess, image_lists, FLAGS.architecture, FLAGS.saved_model_dir)
1239
1240
1241if __name__ == '__main__':
1242 parser = argparse.ArgumentParser()
1243 parser.add_argument(
1244 '--image_dir',
1245 type=str,
1246 default='',
1247 help='Path to folders of labeled images.'
1248 )
1249 parser.add_argument(
1250 '--output_graph',
1251 type=str,
1252 default='/tmp/output_graph.pb',
1253 help='Where to save the trained graph.'
1254 )
1255 parser.add_argument(
1256 '--intermediate_output_graphs_dir',
1257 type=str,
1258 default='/tmp/intermediate_graph/',
1259 help='Where to save the intermediate graphs.'
1260 )
1261 parser.add_argument(
1262 '--intermediate_store_frequency',
1263 type=int,
1264 default=0,
1265 help="""\
1266 How many steps to store intermediate graph. If "0" then will not
1267 store.\
1268 """
1269 )
1270 parser.add_argument(
1271 '--output_labels',
1272 type=str,
1273 default='/tmp/output_labels.txt',
1274 help='Where to save the trained graph\'s labels.'
1275 )
1276 parser.add_argument(
1277 '--summaries_dir',
1278 type=str,
1279 default='/tmp/retrain_logs',
1280 help='Where to save summary logs for TensorBoard.'
1281 )
1282 parser.add_argument(
1283 '--how_many_training_steps',
1284 type=int,
1285 default=4000,
1286 help='How many training steps to run before ending.'
1287 )
1288 parser.add_argument(
1289 '--learning_rate',
1290 type=float,
1291 default=0.01,
1292 help='How large a learning rate to use when training.'
1293 )
1294 parser.add_argument(
1295 '--testing_percentage',
1296 type=int,
1297 default=10,
1298 help='What percentage of images to use as a test set.'
1299 )
1300 parser.add_argument(
1301 '--validation_percentage',
1302 type=int,
1303 default=10,
1304 help='What percentage of images to use as a validation set.'
1305 )
1306 parser.add_argument(
1307 '--eval_step_interval',
1308 type=int,
1309 default=10,
1310 help='How often to evaluate the training results.'
1311 )
1312 parser.add_argument(
1313 '--train_batch_size',
1314 type=int,
1315 default=100,
1316 help='How many images to train on at a time.'
1317 )
1318 parser.add_argument(
1319 '--test_batch_size',
1320 type=int,
1321 default=-1,
1322 help="""\
1323 How many images to test on. This test set is only used once, to evaluate
1324 the final accuracy of the model after training completes.
1325 A value of -1 causes the entire test set to be used, which leads to more
1326 stable results across runs.\
1327 """
1328 )
1329 parser.add_argument(
1330 '--validation_batch_size',
1331 type=int,
1332 default=100,
1333 help="""\
1334 How many images to use in an evaluation batch. This validation set is
1335 used much more often than the test set, and is an early indicator of how
1336 accurate the model is during training.
1337 A value of -1 causes the entire validation set to be used, which leads to
1338 more stable results across training iterations, but may be slower on large
1339 training sets.\
1340 """
1341 )
1342 parser.add_argument(
1343 '--print_misclassified_test_images',
1344 default=False,
1345 help="""\
1346 Whether to print out a list of all misclassified test images.\
1347 """,
1348 action='store_true'
1349 )
1350 parser.add_argument(
1351 '--model_dir',
1352 type=str,
1353 default='/tmp/imagenet',
1354 help="""\
1355 Path to classify_image_graph_def.pb,
1356 imagenet_synset_to_human_label_map.txt, and
1357 imagenet_2012_challenge_label_map_proto.pbtxt.\
1358 """
1359 )
1360 parser.add_argument(
1361 '--bottleneck_dir',
1362 type=str,
1363 default='/tmp/bottleneck',
1364 help='Path to cache bottleneck layer values as files.'
1365 )
1366 parser.add_argument(
1367 '--final_tensor_name',
1368 type=str,
1369 default='final_result',
1370 help="""\
1371 The name of the output classification layer in the retrained graph.\
1372 """
1373 )
1374 parser.add_argument(
1375 '--flip_left_right',
1376 default=False,
1377 help="""\
1378 Whether to randomly flip half of the training images horizontally.\
1379 """,
1380 action='store_true'
1381 )
1382 parser.add_argument(
1383 '--random_crop',
1384 type=int,
1385 default=0,
1386 help="""\
1387 A percentage determining how much of a margin to randomly crop off the
1388 training images.\
1389 """
1390 )
1391 parser.add_argument(
1392 '--random_scale',
1393 type=int,
1394 default=0,
1395 help="""\
1396 A percentage determining how much to randomly scale up the size of the
1397 training images by.\
1398 """
1399 )
1400 parser.add_argument(
1401 '--random_brightness',
1402 type=int,
1403 default=0,
1404 help="""\
1405 A percentage determining how much to randomly multiply the training image
1406 input pixels up or down by.\
1407 """
1408 )
1409 parser.add_argument(
1410 '--architecture',
1411 type=str,
1412 default='inception_v3',
1413 help="""\
1414 Which model architecture to use. 'inception_v3' is the most accurate, but
1415 also the slowest. For faster or smaller models, chose a MobileNet with the
1416 form 'mobilenet_<parameter size>_<input_size>[_quantized]'. For example,
1417 'mobilenet_1.0_224' will pick a model that is 17 MB in size and takes 224
1418 pixel input images, while 'mobilenet_0.25_128_quantized' will choose a much
1419 less accurate, but smaller and faster network that's 920 KB on disk and
1420 takes 128x128 images. See https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html
1421 for more information on Mobilenet.\
1422 """)
1423 parser.add_argument(
1424 '--saved_model_dir',
1425 type=str,
1426 default='/tmp/saved_models/1/',
1427 help='Where to save the exported graph.'
1428 )
1429 FLAGS, unparsed = parser.parse_known_args()
1430 tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)