def test_independent_sample_hierarchical(self):
structure = schema.OneOf([
schema.OneOf(['a', 'b', 'c'], basic_specs.OP_TAG),
schema.OneOf(['d', 'e', 'f', 'g'], basic_specs.OP_TAG),
], basic_specs.OP_TAG)
rl_structure, dist_info = controller.independent_sample(
structure,
increase_ops_probability=0,
increase_filters_probability=0,
hierarchical=True)
tensors = {
'outer_mask': rl_structure.mask,
'entropy': dist_info['entropy'],
'sample_log_prob': dist_info['sample_log_prob'],
}
self.evaluate(tf.global_variables_initializer())
for _ in range(10):
values = self.evaluate(tensors)
if np.all(values['outer_mask'] == np.array([1, 0])):
self.assertAlmostEqual(values['entropy'],
math.log(2) + math.log(3))
self.assertAlmostEqual(values['sample_log_prob'],
math.log(1 / 2) + math.log(1 / 3))
elif np.all(values['outer_mask'] == np.array([0, 1])):
self.assertAlmostEqual(values['entropy'],
math.log(2) + math.log(4))
self.assertAlmostEqual(values['sample_log_prob'],
math.log(1 / 2) + math.log(1 / 4))
else:
self.fail('Unexpected outer_mask: %s', values['outer_mask'])
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 _build_model(params,
model_spec):
"""Translate a ConvTowerSpec namedtuple into a rematlib Layer."""
input_filters = schema.OneOf([3], basic_specs.FILTERS_TAG)
layer = None
result = []
endpoints = []
for block_spec in model_spec.blocks:
for layer_spec in block_spec.layers:
if isinstance(layer_spec, mobile_search_space_v3.DetectionEndpointSpec):
if layer is None:
raise ValueError(
'The first layer of the network cannot be a detection endpoint.')
endpoints.append(layer)
else:
output_filters = block_spec.filters
if isinstance(output_filters, int):
output_filters = schema.OneOf(
[output_filters], basic_specs.FILTERS_TAG)
layer = _build_layer(
params, layer_spec, input_filters, output_filters,
model_spec.filters_base)
input_filters = output_filters
result.append(layer)
result.append(_make_head(params))
# Build the model
model = layers.Sequential(result, aux_outputs=endpoints)
return model
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_tf_indices_with_masks(self):
model_spec = [
schema.OneOf([1], 'foo', mask=tf.constant([1])),
schema.OneOf([2, 3], 'bar', mask=tf.constant([0, 1])),
schema.OneOf([4, 5, 6], 'baz', mask=tf.constant([0, 0, 1])),
]
indices = search_space_utils.tf_indices(model_spec)
self.assertAllEqual(self.evaluate(indices), [0, 1, 2])
def test_get_mask(self):
with self.assertRaisesWithPredicateMatch(ValueError,
'OneOf must have a mask'):
cost_model_lib.get_mask(
schema.OneOf([6, 7, 8, 9], basic_specs.OP_TAG))
mask = cost_model_lib.get_mask(
schema.OneOf([6, 7, 8, 9], basic_specs.OP_TAG,
tf.constant([0, 0, 1, 0])))
self.assertAllEqual(self.evaluate(mask), [0, 0, 1, 0])
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_simple_model_spec_validation_no_tags(self):
model_spec = [
schema.OneOf(['a', 'b', 'c'], basic_specs.OP_TAG),
schema.OneOf(['z', 'w'], basic_specs.FILTERS_TAG),
]
with self.assertRaisesWithLiteralMatch(
ValueError,
'Genotype contains 1 oneofs but model_spec contains 2'):
search_space_utils.prune_model_spec(model_spec=model_spec,
genotype=[0],
prune_filters_by_value=True)
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_prune_filters_by_value_with_invalid_value(self):
model_spec = {
'op': schema.OneOf(['a', 'b', 'c'], basic_specs.OP_TAG),
'filters': schema.OneOf([128, 256, 512], basic_specs.FILTERS_TAG),
}
genotype = {
basic_specs.OP_TAG: [1],
basic_specs.FILTERS_TAG: [1024], # Use values instead of indices
}
with self.assertRaises(ValueError):
search_space_utils.prune_model_spec(model_spec,
genotype,
prune_filters_by_value=True)
def test_prune_simple_model_spec_no_tags(self):
pruned_spec = search_space_utils.prune_model_spec(
model_spec=[
schema.OneOf(['a', 'b', 'c'], basic_specs.OP_TAG),
schema.OneOf(['z', 'w'], basic_specs.FILTERS_TAG),
],
genotype=[2, 1],
prune_filters_by_value=True)
self.assertEqual(pruned_spec, [
schema.OneOf(['c'], basic_specs.OP_TAG),
schema.OneOf(['w'], basic_specs.FILTERS_TAG),
])
def test_map_oenofs_with_tuple_paths_trivial(self):
structure = schema.OneOf([1, 2], 'tag')
all_paths = []
all_oneofs = []
def visit(path, oneof):
all_paths.append(path)
all_oneofs.append(oneof)
return schema.OneOf([x * 10 for x in oneof.choices], oneof.tag)
self.assertEqual(schema.map_oneofs_with_tuple_paths(visit, structure),
schema.OneOf([10, 20], 'tag'))
self.assertEqual(all_paths, [()])
self.assertEqual(all_oneofs, [schema.OneOf([1, 2], 'tag')])
def _from_json(structure):
"""Converted a pure JSON data structure to one with namedtuples and OneOfs."""
if structure is None or isinstance(structure, _PRIMITIVE_TYPES):
return structure
elif isinstance(structure, list):
assert structure
typename = structure[0]
structure = structure[1:]
if typename == 'dict':
return {_from_json(k): _from_json(v) for (k, v) in structure}
elif typename.startswith('namedtuple:'):
cls = namedtuple_name_to_class(typename[len('namedtuple:'):])
kv_pairs = [(_from_json(k), _from_json(v)) for (k, v) in structure]
return _namedtuple_from_json(cls, kv_pairs)
elif typename == 'oneof':
keys = tuple(_from_json(k) for (k, v) in structure)
assert keys == ('choices', 'tag'), keys
return schema.OneOf(*(_from_json(v) for (k, v) in structure))
elif typename == 'list':
return list(map(_from_json, structure))
elif typename == 'tuple':
return tuple(map(_from_json, structure))
else:
raise ValueError('Unsupported __type: {}'.format(typename))
else:
raise ValueError('Unrecognized JSON type: {}'.format(type(structure)))
def test_independent_sample_temperature(self):
structure = schema.OneOf(['foo', 'bar', 'baz'], basic_specs.OP_TAG)
temperature = tf.placeholder_with_default(tf.constant(5.0, tf.float32),
shape=(),
name='temperature')
rl_structure, dist_info = controller.independent_sample(
structure, temperature=temperature)
with self.cached_session() as sess:
sess.run(tf.global_variables_initializer())
# Samples should be valid even when the temperature is set to a value
# other than 1.
self.assertOneHot(sess.run(rl_structure.mask))
# Before training, the sample log-probability and entropy shouldn't be
# affected by the temperature, since the probabilities are initialized
# to a uniform distribution.
self.assertAlmostEqual(sess.run(dist_info['sample_log_prob']),
math.log(1 / 3))
self.assertAlmostEqual(sess.run(dist_info['entropy']), math.log(3))
# The gradients should be multiplied by (1 / temperature).
# The OneOf has three possible choices. The gradient for the selected one
# will be positive, while the gradients for the other two will be
# negative. Since the selected choice can change between steps, we compare
# the max, which should always give us gradients w.r.t. the selected one.
trainable_vars = tf.trainable_variables()
self.assertLen(trainable_vars, 1)
grad_tensors = tf.gradients(dist_info['sample_log_prob'],
trainable_vars)
grad1 = np.max(sess.run(grad_tensors[0], {temperature: 1.0}))
grad5 = np.max(sess.run(grad_tensors[0], {temperature: 5.0}))
self.assertAlmostEqual(grad1 / 5, grad5)
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_tf_argmax_or_zero(self):
# Pruned OneOf structure without mask
oneof = schema.OneOf([1], 'foo')
self.assertEqual(
self.evaluate(search_space_utils.tf_argmax_or_zero(oneof)), 0)
# Unpruned OneOf structure without mask
with self.assertRaisesWithPredicateMatch(ValueError,
'Expect pruned structure'):
oneof = schema.OneOf([1, 2], 'foo')
search_space_utils.tf_argmax_or_zero(oneof)
# Unpruned OneOf structure with mask
oneof = schema.OneOf([1, 2], 'foo', tf.constant([0.0, 1.0]))
self.assertEqual(
self.evaluate(search_space_utils.tf_argmax_or_zero(oneof)), 1)
def test_serialization_with_simple_structures(self):
# Primitives.
self._run_serialization_test(None)
self._run_serialization_test(1)
self._run_serialization_test(0.5)
self._run_serialization_test(1.0)
self._run_serialization_test('foo')
# Lists and tuples.
self._run_serialization_test([1, 2, 3])
self._run_serialization_test((1, 2, 3))
# Dictionaries.
self._run_serialization_test({'a': 3, 'b': 4})
self._run_serialization_test({10: 'x', 20: 'y'})
self._run_serialization_test({(1, 2): 'x', (3, 4): 'y'})
# Namedtuples
self._run_serialization_test(NamedTuple1(42),
expected_type=NamedTuple1)
self._run_serialization_test(NamedTuple2(12345),
expected_type=NamedTuple2)
# OneOf nodes.
self._run_serialization_test(schema.OneOf((1, 2, 3), 'tag'))
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)
def test_prune_model_spec_prune_filters_by_value(self):
model_spec = {
'op': schema.OneOf(['a', 'b', 'c'], basic_specs.OP_TAG),
'filters': schema.OneOf([128, 256, 512], basic_specs.FILTERS_TAG),
}
genotype = {
basic_specs.OP_TAG: [1],
basic_specs.FILTERS_TAG: [512], # Use values instead of indices
}
pruned_spec = search_space_utils.prune_model_spec(
model_spec, genotype, prune_filters_by_value=True)
expected_spec = {
'op': schema.OneOf(['b'], basic_specs.OP_TAG),
'filters': schema.OneOf([512], basic_specs.FILTERS_TAG),
}
self.assertEqual(pruned_spec, expected_spec)
def test_independent_sample_increase_ops_probability_1(self):
structure = schema.OneOf(['foo', 'bar', 'baz'], basic_specs.OP_TAG)
rl_structure, dist_info = controller.independent_sample(
structure, increase_ops_probability=1.0)
self.evaluate(tf.global_variables_initializer())
self.assertAllClose(self.evaluate(rl_structure.mask), [1/3, 1/3, 1/3])
self.assertEqual(self.evaluate(dist_info['sample_log_prob']), 0)
def sepconv(s):
choices = []
for kernel_size in (3, 5, 7):
choices.append(
SeparableConvSpec(kernel_size=kernel_size,
strides=s,
activation=RELU))
return schema.OneOf(choices, 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
def test_independent_sample_increase_ops_does_not_affect_ops(self):
structure = schema.OneOf([42, 64], basic_specs.OP_TAG)
rl_structure, dist_info = controller.independent_sample(
structure, increase_filters_probability=1.0)
self.evaluate(tf.global_variables_initializer())
self.assertOneHot(self.evaluate(rl_structure.mask))
self.assertAlmostEqual(self.evaluate(dist_info['sample_log_prob']),
math.log(1 / 2))
def test_independent_sample_increase_filters_probability_0(self):
structure = schema.OneOf([4, 12, 8], basic_specs.FILTERS_TAG)
rl_structure, dist_info = controller.independent_sample(
structure, increase_filters_probability=0.0)
self.evaluate(tf.global_variables_initializer())
self.assertOneHot(self.evaluate(rl_structure.mask))
self.assertAlmostEqual(self.evaluate(dist_info['sample_log_prob']),
math.log(1 / 3))
def test_independent_sample_increase_filters_probability_1_big_space(self):
# Use a large enough number of choices that we're unlikely to select the
# right one by random chance.
structure = schema.OneOf(list(range(100)), basic_specs.FILTERS_TAG)
rl_structure, dist_info = controller.independent_sample(
structure, increase_filters_probability=1.0)
self.evaluate(tf.global_variables_initializer())
self.assertAllClose(self.evaluate(rl_structure.mask), [0] * 99 + [1])
self.assertEqual(self.evaluate(dist_info['sample_log_prob']), 0)
def test_independent_sample_increase_filters_probability_1(self):
# Make sure that increase_filters does the right thing when the choices do
# not appear in sorted order.
structure = schema.OneOf([4, 12, 8], basic_specs.FILTERS_TAG)
rl_structure, dist_info = controller.independent_sample(
structure, increase_filters_probability=1.0)
self.evaluate(tf.global_variables_initializer())
self.assertAllClose(self.evaluate(rl_structure.mask), [0, 1, 0])
self.assertEqual(self.evaluate(dist_info['sample_log_prob']), 0)
def bneck(input_size, se, s, act):
"""Construct a DepthwiseBottleneckSpec namedtuple."""
if use_relative_expansion_filters:
expansion_filters = sorted({
basic_specs.FilterMultiplier(expansion)
for expansion in expansion_multipliers
})
else:
expansion_filters = sorted({
search_space_utils.scale_filters(input_size, expansion, base=8)
for expansion in expansion_multipliers
})
if search_squeeze_and_excite:
# Replace the default value of the argument 'se' with a OneOf node.
se = schema.OneOf([False, True], basic_specs.OP_TAG)
return DepthwiseBottleneckSpec(
kernel_size=schema.OneOf([3, 5, 7], basic_specs.OP_TAG),
expansion_filters=choose_filters(expansion_filters),
use_squeeze_and_excite=se,
strides=s,
activation=act)
def update_mask(path, oneof):
"""Add a one-hot mask to 'oneof' whose value is derived from 'indices'."""
index = indices[State.position]
State.position += 1
if index < 0 or index >= len(oneof.choices):
raise ValueError(
'Invalid index: {:d} for path: {:s} with {:d} choices'.format(
index, path, len(oneof.choices)))
mask = tf.one_hot(index, len(oneof.choices))
return schema.OneOf(oneof.choices, oneof.tag, mask)
def test_map_oneofs_with_tuple_paths_containing_arrays_and_dicts(self):
structure = {
'foo': [
schema.OneOf([1, 2], 'tag1'),
schema.OneOf([3, 4, 5], 'tag2'),
]
}
all_paths = []
all_oneofs = []
def visit(path, oneof):
all_paths.append(path)
all_oneofs.append(oneof)
return schema.OneOf([x * 10 for x in oneof.choices], oneof.tag)
self.assertEqual(
schema.map_oneofs_with_tuple_paths(visit, structure), {
'foo': [
schema.OneOf([10, 20], 'tag1'),
schema.OneOf([30, 40, 50], 'tag2'),
]
})
self.assertEqual(all_paths, [
('foo', 0),
('foo', 1),
])
self.assertEqual(all_oneofs, [
schema.OneOf([1, 2], 'tag1'),
schema.OneOf([3, 4, 5], 'tag2'),
])
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'),
])
