def test_v3_zero_or_zero(self):
mask = tf.placeholder(shape=[2], dtype=tf.float32)
layer = schema.OneOf(choices=[
basic_specs.ZeroSpec(),
basic_specs.ZeroSpec(),
],
tag=basic_specs.OP_TAG,
mask=mask)
features = mobile_cost_model.coupled_tf_features(
_make_single_layer_model(layer))
with self.session() as sess:
self.assertAllClose([1.0, 0.0],
sess.run(features, {mask: [1.0, 0.0]}))
self.assertAllClose([0.0, 1.0],
sess.run(features, {mask: [0.0, 1.0]}))
def bneck(s, skippable):
"""Construct a spec for an inverted bottleneck layer."""
possible_filter_multipliers = [3.0, 6.0]
possible_kernel_sizes = [3, 5, 7]
choices = []
if collapse_shared_ops:
kernel_size = schema.OneOf(possible_kernel_sizes,
basic_specs.OP_TAG)
expansion_filters = schema.OneOf([
basic_specs.FilterMultiplier(multiplier)
for multiplier in possible_filter_multipliers
], basic_specs.FILTERS_TAG)
choices.append(
DepthwiseBottleneckSpec(kernel_size=kernel_size,
expansion_filters=expansion_filters,
use_squeeze_and_excite=False,
strides=s,
activation=RELU))
else:
for multiplier in possible_filter_multipliers:
for kernel_size in possible_kernel_sizes:
choices.append(
DepthwiseBottleneckSpec(
kernel_size=kernel_size,
expansion_filters=basic_specs.FilterMultiplier(
multiplier),
use_squeeze_and_excite=False,
strides=s,
activation=RELU))
if skippable:
choices.append(basic_specs.ZeroSpec())
return schema.OneOf(choices, basic_specs.OP_TAG)
def test_v3_single_choice_zero_only(self):
layer = schema.OneOf(choices=[basic_specs.ZeroSpec()],
tag=basic_specs.OP_TAG,
mask=tf.constant([1.0]))
features = mobile_cost_model.coupled_tf_features(
_make_single_layer_model(layer))
self.assertAllClose(self.evaluate(features), [1.0])
def test_prune_model_spec_with_path_dropout_rate_tensor(self):
model_spec = {
'op1':
schema.OneOf([
mobile_search_space_v3.ConvSpec(kernel_size=2, strides=2),
basic_specs.ZeroSpec(),
], basic_specs.OP_TAG),
'op2':
schema.OneOf([
mobile_search_space_v3.ConvSpec(kernel_size=3, strides=4),
], basic_specs.OP_TAG),
'filter':
schema.OneOf([32], basic_specs.FILTERS_TAG),
}
model_spec = search_space_utils.prune_model_spec(
model_spec, {basic_specs.OP_TAG: [0, 0]},
path_dropout_rate=tf.constant(2.0) / tf.constant(10.0),
training=True)
self.assertCountEqual(model_spec.keys(), ['op1', 'op2', 'filter'])
self.assertEqual(model_spec['op1'].mask.shape, tf.TensorShape([1]))
self.assertIsNone(model_spec['op2'].mask)
self.assertIsNone(model_spec['filter'].mask)
# The value should either be 0 or 1 / (1 - path_dropout_rate) = 1.25
op_mask_value = self.evaluate(model_spec['op1'].mask)
self.assertTrue(
abs(op_mask_value - 0) < 1e-6 or abs(op_mask_value - 1.25) < 1e-6,
msg='Unexpected op_mask_value: {}'.format(op_mask_value))
def test_get_strides_residual_connection(self):
self.assertEqual(
(1, 1),
mobile_search_space_v3.get_strides(
mobile_search_space_v3.ResidualSpec(basic_specs.ZeroSpec())))
with self.assertRaisesRegex(ValueError,
'Residual layer must have stride 1'):
mobile_search_space_v3.get_strides(
mobile_search_space_v3.ResidualSpec(
mobile_search_space_v3.GlobalAveragePoolSpec()))
def test_v3_zero_or_conv_with_child(self):
kernel_size_mask = tf.placeholder(shape=[3], dtype=tf.float32)
kernel_size = schema.OneOf([3, 5, 7], basic_specs.OP_TAG,
kernel_size_mask)
layer_mask = tf.placeholder(shape=[2], dtype=tf.float32)
layer = schema.OneOf(choices=[
basic_specs.ZeroSpec(),
mobile_search_space_v3.ConvSpec(kernel_size=kernel_size,
strides=1),
],
tag=basic_specs.OP_TAG,
mask=layer_mask)
features = mobile_cost_model.coupled_tf_features(
_make_single_layer_model(layer))
with self.session() as sess:
self.assertAllClose([1.0, 0.0, 0.0, 0.0],
sess.run(
features, {
layer_mask: [1, 0],
kernel_size_mask: [1, 0, 0]
}))
self.assertAllClose([1.0, 0.0, 0.0, 0.0],
sess.run(
features, {
layer_mask: [1, 0],
kernel_size_mask: [0, 1, 0]
}))
self.assertAllClose([1.0, 0.0, 0.0, 0.0],
sess.run(
features, {
layer_mask: [1, 0],
kernel_size_mask: [0, 0, 1]
}))
self.assertAllClose([0.0, 1.0, 0.0, 0.0],
sess.run(
features, {
layer_mask: [0, 1],
kernel_size_mask: [1, 0, 0]
}))
self.assertAllClose([0.0, 0.0, 1.0, 0.0],
sess.run(
features, {
layer_mask: [0, 1],
kernel_size_mask: [0, 1, 0]
}))
self.assertAllClose([0.0, 0.0, 0.0, 1.0],
sess.run(
features, {
layer_mask: [0, 1],
kernel_size_mask: [0, 0, 1]
}))
def test_prune_model_spec_with_path_dropout_training(self):
model_spec = {
'op1':
schema.OneOf([
mobile_search_space_v3.ConvSpec(kernel_size=2, strides=2),
basic_specs.ZeroSpec(),
], basic_specs.OP_TAG),
'op2':
schema.OneOf([
mobile_search_space_v3.ConvSpec(kernel_size=3, strides=4),
], basic_specs.OP_TAG),
'filter':
schema.OneOf([32], basic_specs.FILTERS_TAG),
}
model_spec = search_space_utils.prune_model_spec(
model_spec, {basic_specs.OP_TAG: [0, 0]},
path_dropout_rate=0.2,
training=True)
self.assertCountEqual(model_spec.keys(), ['op1', 'op2', 'filter'])
self.assertEqual(model_spec['op1'].mask.shape, tf.TensorShape([1]))
self.assertIsNone(model_spec['op2'].mask)
self.assertIsNone(model_spec['filter'].mask)
self.assertEqual(
model_spec['op1'].choices,
[mobile_search_space_v3.ConvSpec(kernel_size=2, strides=2)])
self.assertEqual(
model_spec['op2'].choices,
[mobile_search_space_v3.ConvSpec(kernel_size=3, strides=4)])
self.assertEqual(model_spec['filter'].choices, [32])
self.assertEqual(model_spec['op1'].tag, basic_specs.OP_TAG)
self.assertEqual(model_spec['op2'].tag, basic_specs.OP_TAG)
self.assertEqual(model_spec['filter'].tag, basic_specs.FILTERS_TAG)
op_mask_sum = 0
for _ in range(100):
# The value should either be 0 or 1 / (1 - path_dropout_rate) = 1.25
op_mask_value = self.evaluate(model_spec['op1'].mask)
self.assertTrue(
abs(op_mask_value - 0) < 1e-6
or abs(op_mask_value - 1.25) < 1e-6,
msg='Unexpected op_mask_value: {}'.format(op_mask_value))
op_mask_sum += op_mask_value[0]
# The probability of this test failing by random chance is roughly 0.002%.
# Our random number generators are deterministically seeded, so the test
# shouldn't be flakey.
self.assertGreaterEqual(op_mask_sum, 75)
self.assertLessEqual(op_mask_sum, 113)
def test_prune_model_spec_with_path_dropout_eval(self):
model_spec = {
'op1':
schema.OneOf([
mobile_search_space_v3.ConvSpec(kernel_size=2, strides=2),
basic_specs.ZeroSpec(),
], basic_specs.OP_TAG),
'op2':
schema.OneOf([
mobile_search_space_v3.ConvSpec(kernel_size=3, strides=4),
], basic_specs.OP_TAG),
'filter':
schema.OneOf([32], basic_specs.FILTERS_TAG),
}
model_spec = search_space_utils.prune_model_spec(
model_spec, {basic_specs.OP_TAG: [0, 0]},
path_dropout_rate=0.2,
training=False)
self.assertCountEqual(model_spec.keys(), ['op1', 'op2', 'filter'])
# Even though path_dropout_rate=0.2, the controller should not populate
# the mask for op1 because we called prune_model_spec() with training=False.
# In other words, path_dropout_rate should only affect the behavior during
# training, not during evaluation.
self.assertIsNone(model_spec['op1'].mask)
self.assertIsNone(model_spec['op2'].mask)
self.assertIsNone(model_spec['filter'].mask)
self.assertEqual(model_spec['op1'].tag, basic_specs.OP_TAG)
self.assertEqual(model_spec['op2'].tag, basic_specs.OP_TAG)
self.assertEqual(model_spec['filter'].tag, basic_specs.FILTERS_TAG)
self.assertEqual(
model_spec['op1'].choices,
[mobile_search_space_v3.ConvSpec(kernel_size=2, strides=2)])
self.assertEqual(
model_spec['op2'].choices,
[mobile_search_space_v3.ConvSpec(kernel_size=3, strides=4)])
self.assertEqual(model_spec['filter'].choices, [32])
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 optional(layer):
return schema.OneOf([layer, basic_specs.ZeroSpec()],
basic_specs.OP_TAG)
def test_get_strides_zero(self):
self.assertEqual(
(1, 1), mobile_search_space_v3.get_strides(basic_specs.ZeroSpec()))
