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)