def _build_network(self, features, labels, mode):
"""Build a network that returns loss and logits from features and labels."""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
is_predict = (mode == tf.estimator.ModeKeys.PREDICT)
steps_per_epoch = float(
NUM_TRAIN_IMAGES) / self.hparams.train_batch_size
num_total_steps = int(steps_per_epoch * self.hparams.num_epochs)
if getattr(self.hparams, 'num_total_steps', None) is None:
self.hparams.add_hparam('num_total_steps', num_total_steps)
else:
self.hparams.set_hparam('num_total_steps', num_total_steps)
hparams = copy.deepcopy(self.hparams)
if not is_training:
hparams.set_hparam('use_aux_head', False)
tf.logging.info('Amoeba net received hparams for {}:\n{}'.format(
'training' if is_training else 'eval', formatted_hparams(hparams)))
logits, end_points = model_builder.build_network(
features, LABEL_CLASSES, is_training, hparams)
if not is_predict:
labels = tf.one_hot(labels, LABEL_CLASSES)
loss = model_builder.build_softmax_loss(
logits,
end_points,
labels,
label_smoothing=hparams.label_smoothing,
add_summary=False)
# Calculate and print the number of trainable parameters in the model
if is_training:
flops = model_builder.compute_flops_per_example(
hparams.train_batch_size)
else:
flops = model_builder.compute_flops_per_example(
hparams.eval_batch_size)
tf.logging.info('number of flops: {}'.format(flops))
self._calc_num_trainable_params()
if is_predict:
return None, logits
return loss, logits
def _module_fn(is_training):
"""A module_fn for use with hub.create_module_spec().
Args:
is_training: a boolean, passed to the config.network_fn.
This is meant to control whether batch norm, dropout etc. are built
in training or inference mode for this graph version.
Raises:
ValueError: if network_fn outputs are not as expected.
"""
# Set up the module input, and attach an ImageModuleInfo about it.
with tf.name_scope('hub_input'):
default_size = (hparams.image_size, ) * 2
image_module_info = hub.ImageModuleInfo()
size_info = image_module_info.default_image_size
size_info.height, size_info.width = default_size
# TODO(b/72731449): Support variable input size.
shape = (None, ) + default_size + (3, )
images = tf.placeholder(dtype=tf.float32,
shape=shape,
name='images')
hub.attach_image_module_info(image_module_info)
# The input is expected to have RGB color values in the range [0,1]
# and gets converted for AmoebaNet to the Inception-style range [-1,+1].
scaled_images = tf.multiply(images, 2.0)
scaled_images = tf.subtract(scaled_images, 1.0)
# Build the net.
logits, end_points = model_builder.build_network(
scaled_images, num_classes, is_training, hparams)
with tf.name_scope('hub_output'):
# Extract the feature_vectors output.
try:
feature_vectors = end_points['global_pool']
except KeyError:
tf.logging.error('Valid keys of end_points are:',
', '.join(end_points))
raise
with tf.name_scope('feature_vector'):
if feature_vectors.shape.ndims != 2:
raise ValueError(
'Wrong rank (expected 2 after squeeze) '
'in feature_vectors:', feature_vectors)
# Extract the logits output (if applicable).
if num_classes:
with tf.name_scope('classification'):
if logits.shape.ndims != 2:
raise ValueError('Wrong rank (expected 2) in logits:',
logits)
# Add named signatures.
hub.add_signature('image_feature_vector', dict(images=images),
dict(end_points, default=feature_vectors))
if num_classes:
hub.add_signature('image_classification', dict(images=images),
dict(end_points, default=logits))
# Add the default signature.
if num_classes:
hub.add_signature('default', dict(images=images),
dict(default=logits))
else:
hub.add_signature('default', dict(images=images),
dict(default=feature_vectors))