Exemplo n.º 1
0
def build_nasnet_large(images, num_classes,
                       is_training=True,
                       final_endpoint=None,
                       config=None,
                       current_step=None):
    """Build NASNet Large model for the ImageNet Dataset."""
    hparams = (large_imagenet_config() if config is None
               else copy.deepcopy(config))
    _update_hparams(hparams, is_training)

    if tf.test.is_gpu_available() and hparams.data_format == 'NHWC':
        tf.logging.info('A GPU is available on the machine, consider using NCHW '
                        'data format for increased speed on GPU.')

    if hparams.data_format == 'NCHW':
        images = tf.transpose(images, [0, 3, 1, 2])

    # Calculate the total number of cells in the network
    # Add 2 for the reduction cells
    total_num_cells = hparams.num_cells + 2
    # If ImageNet, then add an additional two for the stem cells
    total_num_cells += 2

    normal_cell = nasnet_utils.NasNetANormalCell(
        hparams.num_conv_filters, hparams.drop_path_keep_prob,
        total_num_cells, hparams.total_training_steps,
        hparams.use_bounded_activation)
    reduction_cell = nasnet_utils.NasNetAReductionCell(
        hparams.num_conv_filters, hparams.drop_path_keep_prob,
        total_num_cells, hparams.total_training_steps,
        hparams.use_bounded_activation)
    with arg_scope([slim.dropout, nasnet_utils.drop_path, slim.batch_norm],
                   is_training=is_training):
        with arg_scope([slim.avg_pool2d,
                        slim.max_pool2d,
                        slim.conv2d,
                        slim.batch_norm,
                        slim.separable_conv2d,
                        nasnet_utils.factorized_reduction,
                        nasnet_utils.global_avg_pool,
                        nasnet_utils.get_channel_index,
                        nasnet_utils.get_channel_dim],
                       data_format=hparams.data_format):
            return _build_nasnet_base(images,
                                      normal_cell=normal_cell,
                                      reduction_cell=reduction_cell,
                                      num_classes=num_classes,
                                      hparams=hparams,
                                      is_training=is_training,
                                      stem_type='imagenet',
                                      final_endpoint=final_endpoint,
                                      current_step=current_step)
    def _extract_box_classifier_features(self, proposal_feature_maps, scope):
        """Extracts second stage box classifier features.

    This function reconstructs the "second half" of the NASNet-A
    network after the part defined in `_extract_proposal_features`.

    Args:
      proposal_feature_maps: A 4-D float tensor with shape
        [batch_size * self.max_num_proposals, crop_height, crop_width, depth]
        representing the feature map cropped to each proposal.
      scope: A scope name.

    Returns:
      proposal_classifier_features: A 4-D float tensor with shape
        [batch_size * self.max_num_proposals, height, width, depth]
        representing box classifier features for each proposal.
    """
        del scope

        # Note that we always feed into 2 layers of equal depth
        # where the first N channels corresponds to previous hidden layer
        # and the second N channels correspond to the final hidden layer.
        hidden_previous, hidden = tf.split(proposal_feature_maps, 2, axis=3)

        # Note that what follows is largely a copy of build_nasnet_large() within
        # nasnet.py. We are copying to minimize code pollution in slim.

        # TODO(shlens,skornblith): Determine the appropriate drop path schedule.
        # For now the schedule is the default (1.0->0.7 over 250,000 train steps).
        hparams = nasnet.large_imagenet_config()
        if not self._is_training:
            hparams.set_hparam('drop_path_keep_prob', 1.0)

        # Calculate the total number of cells in the network
        # -- Add 2 for the reduction cells.
        total_num_cells = hparams.num_cells + 2
        # -- And add 2 for the stem cells for ImageNet training.
        total_num_cells += 2

        normal_cell = nasnet_utils.NasNetANormalCell(
            hparams.num_conv_filters, hparams.drop_path_keep_prob,
            total_num_cells, hparams.total_training_steps)
        reduction_cell = nasnet_utils.NasNetAReductionCell(
            hparams.num_conv_filters, hparams.drop_path_keep_prob,
            total_num_cells, hparams.total_training_steps)
        with arg_scope([slim.dropout, nasnet_utils.drop_path],
                       is_training=self._is_training):
            with arg_scope([slim.batch_norm],
                           is_training=self._train_batch_norm):
                with arg_scope([
                        slim.avg_pool2d, slim.max_pool2d, slim.conv2d,
                        slim.batch_norm, slim.separable_conv2d,
                        nasnet_utils.factorized_reduction,
                        nasnet_utils.global_avg_pool,
                        nasnet_utils.get_channel_index,
                        nasnet_utils.get_channel_dim
                ],
                               data_format=hparams.data_format):

                    # This corresponds to the cell number just past 'Cell_11' used by
                    # by _extract_proposal_features().
                    start_cell_num = 12
                    # Note that this number equals:
                    #  start_cell_num + 2 stem cells + 1 reduction cell
                    true_cell_num = 15

                    with slim.arg_scope(nasnet.nasnet_large_arg_scope()):
                        net = _build_nasnet_base(hidden_previous,
                                                 hidden,
                                                 normal_cell=normal_cell,
                                                 reduction_cell=reduction_cell,
                                                 hparams=hparams,
                                                 true_cell_num=true_cell_num,
                                                 start_cell_num=start_cell_num)

        proposal_classifier_features = net
        return proposal_classifier_features