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'), ])