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_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 test_map_oneofs_with_tuple_paths_simple(self):
structure = [
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),
[
schema.OneOf([10, 20], 'tag1'),
schema.OneOf([30, 40, 50], 'tag2'),
])
self.assertEqual(all_paths, [(0, ), (1, )])
self.assertEqual(all_oneofs, [
schema.OneOf([1, 2], 'tag1'),
schema.OneOf([3, 4, 5], 'tag2'),
])
def independent_sample(
structure,
increase_ops_probability = None,
increase_filters_probability = None,
hierarchical = True,
name = None,
temperature = 1.0):
"""Generate a search space specification for an RL controller model.
Each OneOf value is sampled independently of every other; hence the name.
Args:
structure: Nested data structure containing OneOf objects to search over.
increase_ops_probability: Scalar float Tensor or None. If not None, we will
randomly enable all possible operations instead of just the selected
operations with this probability.
increase_filters_probability: Scalar float Tensor or None. If not None, we
will randomly use the largest possible filter sizes with this
probability.
hierarchical: Boolean. If true, the values of the outputs `sample_log_prob`
and `entropy` will only take into account subgraphs that are enabled at
the current training step.
name: Optional name for the newly created TensorFlow scope.
temperature: Positive scalar controlling the temperature to use when
sampling from the RL controller.
Returns:
A tuple (new_structure, dist_info) where `new_structure` is a copy
of `structure` annotated with mask tensors, and `dist_info` is a
dictionary containing information about the sampling distribution
which contains the following keys:
- entropy: Scalar float Tensor, entropy of the current probability
distribution.
- logits_by_path: OrderedDict of rank-1 Tensors, sample-independent logits
for each OneOf in `structure`. Names are derived from OneOf paths.
- logits_by_tag: OrderedDict of rank-1 Tensors, sample-independent logits
for each OneOf in `structure`. Names are derived from OneOf tags.
- sample_log_prob: Scalar float Tensor, log-probability of the current
sample associated with `new_structure`.
"""
with tf.variable_scope(name, 'independent_sample'):
temperature = tf.convert_to_tensor(temperature, tf.float32)
dist_info = {
'entropy': tf.constant(0, tf.float32),
'logits_by_path': collections.OrderedDict(),
'logits_by_tag': collections.OrderedDict(),
'sample_log_prob': tf.constant(0, tf.float32),
}
tag_counters = collections.Counter()
entropies = dict()
log_probs = dict()
is_active = dict()
def visit(tuple_path, oneof):
"""Visit a OneOf node in `structure`."""
string_path = '/'.join(map(str, tuple_path))
num_choices = len(oneof.choices)
logits = tf.get_variable(
name='logits/' + string_path,
initializer=tf.initializers.zeros(),
shape=[num_choices],
dtype=tf.float32)
logits = logits / temperature
tag_name = '{:s}_{:d}'.format(oneof.tag, tag_counters[oneof.tag])
tag_counters[oneof.tag] += 1
dist_info['logits_by_path'][string_path] = logits
dist_info['logits_by_tag'][tag_name] = logits
dist = tfp.distributions.OneHotCategorical(
logits=logits, dtype=tf.float32)
entropies[tuple_path] = dist.entropy()
sample_mask = dist.sample()
sample_log_prob = dist.log_prob(sample_mask)
if oneof.tag == basic_specs.OP_TAG:
sample_mask, sample_log_prob = _replace_sample_with_probability(
sample_mask, sample_log_prob, increase_ops_probability,
tf.constant([1.0/num_choices]*num_choices, tf.float32))
elif oneof.tag == basic_specs.FILTERS_TAG:
# NOTE: While np.argmax() was originally designed to work with integer
# filter sizes, it will also work with any object type that supports
# "less than" and "greater than" operations.
sample_mask, sample_log_prob = _replace_sample_with_probability(
sample_mask, sample_log_prob, increase_filters_probability,
tf.one_hot(np.argmax(oneof.choices), len(oneof.choices)))
log_probs[tuple_path] = sample_log_prob
for i in range(len(oneof.choices)):
tuple_subpath = tuple_path + ('choices', i)
is_active[tuple_subpath] = tf.greater(tf.abs(sample_mask[i]), 1e-6)
return schema.OneOf(choices=oneof.choices,
tag=oneof.tag,
mask=sample_mask)
new_structure = schema.map_oneofs_with_tuple_paths(visit, structure)
assert six.viewkeys(entropies) == six.viewkeys(log_probs)
for path in entropies:
path_is_active = tf.constant(True)
if hierarchical:
for i in range(len(path) + 1):
if path[:i] in is_active:
path_is_active = tf.logical_and(path_is_active, is_active[path[:i]])
path_is_active = tf.cast(path_is_active, tf.float32)
dist_info['entropy'] += entropies[path] * path_is_active
dist_info['sample_log_prob'] += log_probs[path] * path_is_active
return (new_structure, dist_info)
def test_map_oneofs_with_tuple_paths_containing_nested_oneofs(self):
structure = {
'root':
schema.OneOf([
schema.OneOf([
{
'leaf': schema.OneOf([1, 10], 'level2')
},
{
'leaf': schema.OneOf([2, 20], 'level2')
},
], 'level1'),
schema.OneOf([
{
'leaf': schema.OneOf([3, 30], 'level2')
},
{
'leaf': schema.OneOf([4, 40], 'level2')
},
{
'leaf': schema.OneOf([5, 50], 'level2')
},
], 'level1')
], 'level0')
}
all_paths = []
all_oneofs = []
def visit(path, oneof):
all_paths.append(path)
all_oneofs.append(oneof)
return schema.OneOf([oneof.choices[0]], oneof.tag)
self.assertEqual(
schema.map_oneofs_with_tuple_paths(visit, structure), {
'root':
schema.OneOf([
schema.OneOf([
{
'leaf': schema.OneOf([1], 'level2')
},
], 'level1'),
], 'level0')
})
self.assertEqual(all_paths, [
('root', 'choices', 0, 'choices', 0, 'leaf'),
('root', 'choices', 0, 'choices', 1, 'leaf'),
('root', 'choices', 0),
('root', 'choices', 1, 'choices', 0, 'leaf'),
('root', 'choices', 1, 'choices', 1, 'leaf'),
('root', 'choices', 1, 'choices', 2, 'leaf'),
('root', 'choices', 1),
('root', ),
])
# A OneOf node's children should already be updated by the time we visit it.
self.assertEqual(all_oneofs, [
schema.OneOf([1, 10], 'level2'),
schema.OneOf([2, 20], 'level2'),
schema.OneOf([
{
'leaf': schema.OneOf([1], 'level2')
},
{
'leaf': schema.OneOf([2], 'level2')
},
], 'level1'),
schema.OneOf([3, 30], 'level2'),
schema.OneOf([4, 40], 'level2'),
schema.OneOf([5, 50], 'level2'),
schema.OneOf([
{
'leaf': schema.OneOf([3], 'level2')
},
{
'leaf': schema.OneOf([4], 'level2')
},
{
'leaf': schema.OneOf([5], 'level2')
},
], 'level1'),
schema.OneOf([
schema.OneOf([
{
'leaf': schema.OneOf([1], 'level2')
},
], 'level1'),
schema.OneOf([
{
'leaf': schema.OneOf([3], 'level2')
},
], 'level1')
], 'level0'),
])
