def test_get_model_spec_with_single_model(self):
model_spec = mobile_classifier_factory.get_model_spec(
mobile_search_space_v3.MOBILENET_V3_LARGE)
self.assertIsInstance(model_spec, basic_specs.ConvTowerSpec)
def validate(oneof):
self.assertLen(oneof.choices, 1)
schema.map_oneofs(validate, model_spec)
def test_mobilenet_v3_like_search_gradients(self):
model_spec = mobile_search_space_v3.mobilenet_v3_like_search()
model_spec = schema.map_oneofs(_with_largest_possible_masks, model_spec)
model = mobile_model_v3.get_model(
model_spec, num_classes=1001, force_stateless_batch_norm=True)
inputs = tf.random_normal(shape=[8, 224, 224, 3], dtype=tf.float32)
model.build(inputs.shape)
logits, unused_endpoints = model.apply(inputs, training=True)
labels = tf.one_hot([0]*8, 1001, dtype=tf.float32)
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits))
variables = model.trainable_variables()
grads = tf.gradients(loss, variables)
for variable, grad in zip(variables, grads):
self.assertIsNotNone(
grad, msg='Gradient for {} is None'.format(variable.name))
self.evaluate(tf.global_variables_initializer())
for variable, array in zip(variables, self.evaluate(grads)):
self.assertFalse(
np.all(np.equal(array, 0)),
msg='Gradient for {} is identically zero'.format(variable.name))
def scale_conv_tower_spec(model_spec, multipliers, base=None):
"""Scale all the filters in `model_spec`, rounding to multiples of `base`.
Args:
model_spec: A ConvTowerSpec namedtuple.
multipliers: float or list/tuple of floats, the possible filter multipliers.
base: Positive integer, all filter sizes must be a multiple of this value.
Returns:
A new basic_specs.ConvTowerSpec.
"""
if base is None:
base = model_spec.filters_base
if isinstance(multipliers, (int, float)):
multipliers = (multipliers, )
def update(oneof):
"""Compute version of `oneof` whose filters have been scaled up/down."""
if oneof.tag != basic_specs.FILTERS_TAG:
return oneof
all_filters = set()
for filters in oneof.choices:
if isinstance(filters, basic_specs.FilterMultiplier):
# Skip scaling because the filter sizes are relative, not absolute.
all_filters.add(filters)
else:
for mult in multipliers:
all_filters.add(scale_filters(filters, mult, base))
return schema.OneOf(sorted(all_filters), basic_specs.FILTERS_TAG)
result = schema.map_oneofs(update, model_spec)
return basic_specs.ConvTowerSpec(result.blocks, base)
def with_random_pruning(model_spec):
"""Pick a random value for each OneOf and prune away the remaining choices."""
def update(oneof):
index = random.randrange(len(oneof.choices))
return schema.OneOf(choices=[oneof.choices[index]], tag=oneof.tag)
return schema.map_oneofs(update, model_spec)
def with_random_masks(model_spec):
"""Assign random one-hot masks OneOf."""
def update(oneof):
mask = random_one_hot(len(oneof.choices))
return schema.OneOf(choices=oneof.choices, tag=oneof.tag, mask=mask)
return schema.map_oneofs(update, model_spec)
def test_map_oneofs(self):
structure = {
'foo': [
schema.OneOf([1, 2], 'tag1'),
schema.OneOf([3, 4, 5], 'tag2'),
]
}
all_oneofs = []
def visit(oneof):
all_oneofs.append(oneof)
return schema.OneOf([x * 10 for x in oneof.choices], oneof.tag)
self.assertEqual(
schema.map_oneofs(visit, structure), {
'foo': [
schema.OneOf([10, 20], 'tag1'),
schema.OneOf([30, 40, 50], 'tag2'),
]
})
self.assertEqual(all_oneofs, [
schema.OneOf([1, 2], 'tag1'),
schema.OneOf([3, 4, 5], 'tag2'),
])
def _validate_genotype_sequence(model_spec,
genotype):
"""Verify that the number of OneOfs in `genotype` matches `model_spec`."""
# Note: Conceptually, we just need oneof_count to be an integer. But we need
# to be able to update its value from within the update_count() function, and
# storing it inside a dictionary makes that easier.
oneof_count = {'value': 0}
def update_count(oneof):
del oneof # Unused
oneof_count['value'] += 1
schema.map_oneofs(update_count, model_spec)
if len(genotype) != oneof_count['value']:
raise ValueError(
'Genotype contains {:d} oneofs but model_spec contains {:d}'
.format(len(genotype), oneof_count['value']))
def tf_indices(model_spec):
"""Extract `indices` from `model_spec` as Tensors.
Args:
model_spec: Nested data structure containing schema.OneOf objects.
Returns:
`indices`, a rank-1 integer Tensor.
"""
indices = []
def visit(oneof):
index = tf_argmax_or_zero(oneof)
indices.append(index)
schema.map_oneofs(visit, model_spec)
return tf.stack(indices)
def with_random_op_masks(model_spec):
"""Assign random one-hot masks OneOf whose tags are basic_specs.OP_TAG."""
def update(oneof):
if oneof.tag == basic_specs.OP_TAG:
mask = random_one_hot(len(oneof.choices))
return schema.OneOf(choices=oneof.choices, tag=oneof.tag, mask=mask)
else:
return oneof
return schema.map_oneofs(update, model_spec)
def _assert_correct_oneof_count(indices,
model_spec):
"""Ensure the length of indices matches the number of OneOfs in model_spec."""
# We use an object with static member fields to maintain internal state so
# that the elements inside can be updated within a nested function.
class State(object):
count = 0 # Total number of oneofs in 'model_spec'
# Count the number of elements in model_spec.
def update_count(oneof):
del oneof # Unused
State.count += 1
schema.map_oneofs(update_count, model_spec)
if State.count != len(indices):
raise ValueError('Wrong number of indices. Expected: {} but got: {}'.format(
State.count, len(indices)))
def _validate_genotype_dict(model_spec,
genotype):
"""Verify that the tag counts in `genotype` match those in `model_spec`."""
# Count the number of times each tag appears in ConvTowerSpec.
tag_counts = collections.Counter()
def update_tag_counts(oneof):
tag_counts[oneof.tag] += 1
schema.map_oneofs(update_tag_counts, model_spec)
# Report any size mismatches we come across.
bad_tags = set(genotype) - set(tag_counts)
if bad_tags:
raise ValueError(
'Tag(s) appear in genotype but not in model_spec: {:s}'
.format(', '.join(bad_tags)))
for tag in genotype:
if len(genotype[tag]) != tag_counts[tag]:
raise ValueError(
'Tag {:s} appears {:d} times in genotype but {:d} times in '
'model_spec'.format(tag, len(genotype[tag]), tag_counts[tag]))
def prune_model_spec(model_spec,
genotype,
path_dropout_rate=0.0,
training=None,
prune_filters_by_value=False):
"""Creates a representation for an architecture with constant ops.
Args:
model_spec: Nested data structure containing schema.OneOf objects.
genotype: A dictionary mapping tags to sequences of integers. Or a sequence
of integers containing the selections for all the OneOf nodes in
model_spec.
path_dropout_rate: Float or scalar float Tensor between 0 and 1. If greater
than zero, we will randomly zero out skippable operations during
training with this probability. Cannot be used with an rl controller.
Should be set to 0 at evaluation time.
training: Boolean. True during training, false during evaluation/inference.
Can be None if path_dropout_rate is zero.
prune_filters_by_value: Boolean. If true, treat genotype[FILTERS_TAG] as a
list of values rather than a list of indices.
Returns:
A pruned version of `model_spec` with all unused options removed.
"""
if path_dropout_rate != 0.0:
if basic_specs.OP_TAG not in genotype:
raise ValueError(
'If path_dropout_rate > 0 then genotype must contain key {:s}.'
.format(basic_specs.OP_TAG))
if training is None:
raise ValueError(
'If path_dropout_rate > 0 then training cannot be None.')
# Create a mutable copy of 'genotype'. This will let us modify the copy
# without updating the original.
genotype_is_dict = isinstance(genotype, dict)
if genotype_is_dict:
genotype = {key: list(value) for (key, value) in genotype.items()}
_validate_genotype_dict(model_spec, genotype)
else: # genotype is a list/tuple of integers
genotype = list(genotype)
_validate_genotype_sequence(model_spec, genotype)
# Everything looks good. Now prune the model.
zero_spec = basic_specs.ZeroSpec()
def update_spec(oneof):
"""Visit a schema.OneOf node in `model_spec`, return an updated value."""
if genotype_is_dict and oneof.tag not in genotype:
return oneof
if genotype_is_dict:
selection = genotype[oneof.tag].pop(0)
if oneof.tag == basic_specs.FILTERS_TAG and prune_filters_by_value:
selection = oneof.choices.index(selection)
else:
selection = genotype.pop(0)
# If an operation is skippable (i.e., it can be replaced with a ZeroSpec)
# then we optionally apply path dropout during stand-alone training.
# This logic, if enabled, will replace a standard RL controller.
mask = None
if (path_dropout_rate != 0.0 and training
and oneof.tag == basic_specs.OP_TAG
and zero_spec in oneof.choices):
keep_prob = 1.0 - path_dropout_rate
# Mask is [1] with probability `keep_prob`, and [0] otherwise.
mask = tf.cast(tf.less(tf.random_uniform([1]), keep_prob),
tf.float32)
# Normalize the mask so that the expected value of each element 1.
mask = mask / keep_prob
return schema.OneOf([oneof.choices[selection]], oneof.tag, mask)
return schema.map_oneofs(update_spec, model_spec)
def test_coupled_tf_features_with_mobile_model_v3(self, ssd):
model_spec = mobile_search_space_v3.get_search_space_spec(ssd)
model_spec = schema.map_oneofs(_assign_random_mask, model_spec)
features = mobile_cost_model.coupled_tf_features(model_spec)
self.assertEqual(features.dtype, tf.float32)
self.assertEqual(features.shape.rank, 1)
