예제 #1
0
    def __init__(self,
                 depth,
                 stride,
                 shortcut_connection,
                 shortcut_from_preact,
                 weight_decay,
                 batch_norm_momentum,
                 batch_norm_epsilon,
                 batch_norm_center,
                 batch_norm_scale,
                 group_norm_groups,
                 name):
        """Constructor.
        Args:
          depth: int scalar, the depth of the two conv ops in each Resnet unit.
          stride: int scalar, the stride of the first conv op in each Resnet unit.
          shortcut_connection: bool scalar, whether to add shortcut connection in
            each Resnet unit. If False, degenerates to a 'Plain network'.
          shortcut_from_preact: bool scalar, whether the shortcut connection starts
            from the preactivation or the input feature map.
          weight_decay: float scalar, weight for l2 regularization.
          batch_norm_momentum: float scalar, the moving average decay.
          batch_norm_epsilon: float scalar, small value to avoid divide by zero.
          batch_norm_center: bool scalar, whether to center in the batch norm.
          batch_norm_scale: bool scalar, whether to scale in the batch norm.
        """
        super(ResNetUnit, self).__init__(name=name)
        self._depth = depth
        self._stride = stride
        self._shortcut_connection = shortcut_connection
        self._shortcut_from_preact = shortcut_from_preact
        self._weight_decay = weight_decay

        self._kernel_regularizer = regularizers.l2(weight_decay)

        self._bn1 = tfalayers.GroupNormalization(groups=group_norm_groups,
                                                 axis=-1,
                                                 epsilon=batch_norm_epsilon,
                                                 center=batch_norm_center,
                                                 scale=batch_norm_scale)
        self._conv1 = layers.Conv2D(depth,
                                    3,
                                    stride,
                                    'same',
                                    use_bias=True,
                                    kernel_regularizer=self._kernel_regularizer,
                                    name='conv1')

        self._bn2 = tfalayers.GroupNormalization(groups=group_norm_groups,
                                                 axis=-1,
                                                 epsilon=batch_norm_epsilon,
                                                 center=batch_norm_center,
                                                 scale=batch_norm_scale)
        self._conv2 = layers.Conv2D(depth,
                                    3,
                                    1,
                                    'same',
                                    use_bias=True,
                                    kernel_regularizer=self._kernel_regularizer,
                                    name='conv2')
예제 #2
0
def normalization(norm_type: str,
                  axis=-1,
                  epsilon=0.001,
                  momentum=0.99,
                  groups=8,
                  name=None):
  """Normalization after conv layers."""
  if norm_type == 'gn':
    return tfa_layers.GroupNormalization(groups, axis, epsilon, name=name)

  if norm_type == 'tpu_bn':
    return TpuBatchNormalization(
        axis=axis, momentum=momentum, epsilon=epsilon, name=name)

  return BatchNormalization(
      axis=axis, momentum=momentum, epsilon=epsilon, name=name)
예제 #3
0
 def __init__(self, hidden_size, projection_size, weight_decay,
              group_norm_groups, batch_norm_epsilon, batch_norm_center,
              batch_norm_scale, use_mlp_norm):
     super().__init__()
     self._kernel_regularizer = tf.keras.regularizers.l2(weight_decay)
     self.dense_in = tf.keras.layers.Dense(
         hidden_size,
         activation='linear',
         kernel_regularizer=self._kernel_regularizer)
     # maybe using groupnorm?
     self.use_mlp_norm = use_mlp_norm
     if self.use_mlp_norm:
         self._bn1 = tfalayers.GroupNormalization(
             groups=group_norm_groups,
             axis=-1,
             epsilon=batch_norm_epsilon,
             center=batch_norm_center,
             scale=batch_norm_scale)
     self.dense1_out = tf.keras.layers.Dense(
         projection_size, kernel_regularizer=self._kernel_regularizer)
예제 #4
0
def inverted_res_block(input_tensor,
                       expansion_factor,
                       stride,
                       filters,
                       alpha,
                       block_number,
                       num_groups=2,
                       dropout_prob=None,
                       expansion_layer=True):
    """Creates an inverted residual block.

  Args:
    input_tensor: A 4D input tensor, with shape (samples, channels, rows, cols)
      or (samples, rows, cols, channels).
    expansion_factor: A positive integer that governs (multiplicatively) how
      many channels are added in the initial expansion layer.
    stride: A positive integer giving the stride of the depthwise convolutional
      layer.
    filters: The base number of filters in the projection layer.
    alpha: A float multiplier for the number of filters in the projection layer.
      If set to 1.0, we use the number of filters is given by the`filters` arg.
    block_number: An integer specifying which inverted residual layer this is.
      Used only for naming purposes.
    num_groups: The number of groups to use in the GroupNorm layers.
    dropout_prob: The probability of setting a weight to zero in the dropout
      layer. If None, no dropout is used.
    expansion_layer: Whether to use an initial expansion layer.

  Returns:
    A 4D tensor with the same shape as the input tensor.
  """

    if tf.keras.backend.image_data_format() == 'channels_last':
        row_axis = 1
        col_axis = 2
        channel_axis = 3
    else:
        channel_axis = 1
        row_axis = 2
        col_axis = 3

    image_shape = (input_tensor.shape[row_axis], input_tensor.shape[col_axis])
    num_input_channels = input_tensor.shape[channel_axis]
    x = input_tensor
    prefix = 'block_{}_'.format(block_number)

    if expansion_layer:
        # We perform an initial pointwise convolution layer.
        x = tf.keras.layers.Conv2D(expansion_factor * num_input_channels,
                                   kernel_size=1,
                                   padding='same',
                                   use_bias=False,
                                   activation=None,
                                   name=prefix + 'expand_conv')(x)
        x = tfa_layers.GroupNormalization(groups=num_groups,
                                          axis=channel_axis,
                                          name=prefix + 'expand_gn')(x)
        if dropout_prob:
            x = tf.keras.layers.Dropout(dropout_prob,
                                        name=prefix + 'expand_dropout')(x)
        x = tf.keras.layers.ReLU(6.0, name=prefix + 'expand_relu')(x)

    # We now use depthwise convolutions
    if stride % 2 == 0:
        padding = compute_pad(image_shape, 3, enforce_odd=True)
        x = tf.keras.layers.ZeroPadding2D(padding=padding,
                                          name=prefix + 'pad')(x)

    padding_type = 'same' if stride == 1 else 'valid'
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=3,
                                        strides=stride,
                                        activation=None,
                                        use_bias=False,
                                        padding=padding_type,
                                        name=prefix + 'depthwise_conv')(x)
    x = tfa_layers.GroupNormalization(groups=num_groups,
                                      axis=channel_axis,
                                      name=prefix + 'depthwise_gn')(x)
    if dropout_prob:
        x = tf.keras.layers.Dropout(dropout_prob,
                                    name=prefix + 'depthwise_dropout')(x)
    x = tf.keras.layers.ReLU(6.0, name=prefix + 'depthwise_relu')(x)

    # Projection phase, using pointwise convolutions
    num_projection_filters = _make_divisible(int(filters * alpha), 8)
    x = tf.keras.layers.Conv2D(num_projection_filters,
                               kernel_size=1,
                               padding='same',
                               use_bias=False,
                               activation=None,
                               name=prefix + 'project_conv')(x)
    x = tfa_layers.GroupNormalization(groups=num_groups,
                                      axis=channel_axis,
                                      name=prefix + 'project_gn')(x)
    if dropout_prob:
        x = tf.keras.layers.Dropout(dropout_prob,
                                    name=prefix + 'project_dropout')(x)
    if num_input_channels == num_projection_filters and stride == 1:
        x = tf.keras.layers.add([input_tensor, x])
    return x
예제 #5
0
def create_small_mobilenet_v2(input_shape,
                              alpha=1.0,
                              pooling='avg',
                              num_groups=2,
                              dropout_prob=None,
                              num_classes=1000):
    """Instantiates a MobileNetV2 model with Group Normalization.

  Args:
    input_shape: A tuple of length 3 describing the number of rows, columns, and
      channels of an input. Can be in channel-first or channel-last format.
    alpha: A float multiplier for the number of filters in the projection
      pointwise convolutional layers. If set to 1.0, we use the default number
      of filters from the original paper.
    pooling: String indicating the pooling mode for the final
      fully-connected layer. Can be one of 'avg' or 'max'.
    num_groups: The number of groups to use in the GroupNorm layers.
    dropout_prob: The probability of setting a weight to zero in the dropout
      layer. If None, no dropout is used.
    num_classes: An integer indicating the number of output classes.

  Returns:
    A `tf.keras.Model`.
  """

    if len(input_shape) != 3:
        raise ValueError('Input shape should be a tuple of length 3.')

    if tf.keras.backend.image_data_format() == 'channels_last':
        row_axis, col_axis = (0, 1)
        channel_axis = 3
    else:
        row_axis, col_axis = (1, 2)
        channel_axis = 1

    image_shape = (input_shape[row_axis], input_shape[col_axis])

    img_input = tf.keras.layers.Input(shape=input_shape)

    initial_padding = compute_pad(image_shape, 3, enforce_odd=True)
    x = tf.keras.layers.ZeroPadding2D(initial_padding,
                                      name='initial_pad')(img_input)

    num_filters_first_block = _make_divisible(32 * alpha, 8)
    x = tf.keras.layers.Conv2D(num_filters_first_block,
                               kernel_size=3,
                               strides=(2, 2),
                               padding='valid',
                               use_bias=False,
                               name='initial_conv')(x)
    x = tfa_layers.GroupNormalization(groups=num_groups,
                                      axis=channel_axis,
                                      name='initial_gn')(x)
    if dropout_prob:
        x = tf.keras.layers.Dropout(dropout_prob, name='initial_dropout')(x)
    x = tf.keras.layers.ReLU(6.0, name='initial_relu')(x)

    x = inverted_res_block(x,
                           expansion_factor=1,
                           stride=1,
                           filters=16,
                           alpha=alpha,
                           block_number=0,
                           num_groups=num_groups,
                           dropout_prob=dropout_prob,
                           expansion_layer=False)

    x = inverted_res_block(x,
                           expansion_factor=6,
                           stride=2,
                           filters=24,
                           alpha=alpha,
                           block_number=1,
                           num_groups=num_groups,
                           dropout_prob=dropout_prob)
    x = inverted_res_block(x,
                           expansion_factor=6,
                           stride=1,
                           filters=24,
                           alpha=alpha,
                           block_number=2,
                           num_groups=num_groups,
                           dropout_prob=dropout_prob)

    x = inverted_res_block(x,
                           expansion_factor=6,
                           stride=2,
                           filters=32,
                           alpha=alpha,
                           block_number=3,
                           num_groups=num_groups,
                           dropout_prob=dropout_prob)
    x = inverted_res_block(x,
                           expansion_factor=6,
                           stride=1,
                           filters=32,
                           alpha=alpha,
                           block_number=4,
                           num_groups=num_groups,
                           dropout_prob=dropout_prob)
    x = inverted_res_block(x,
                           expansion_factor=6,
                           stride=1,
                           filters=32,
                           alpha=alpha,
                           block_number=5,
                           num_groups=num_groups,
                           dropout_prob=dropout_prob)

    # For the last layer, we do not use alpha < 1. This is to recreate the
    # non-usage of alpha in the last layer, as stated in the paper.
    if alpha > 1.0:
        last_block_filters = _make_divisible(640 * alpha, 8)
    else:
        last_block_filters = 640

    x = tf.keras.layers.Conv2D(last_block_filters,
                               kernel_size=1,
                               use_bias=False,
                               name='last_conv')(x)
    x = tfa_layers.GroupNormalization(groups=num_groups,
                                      axis=channel_axis,
                                      name='last_gn')(x)
    if dropout_prob:
        x = tf.keras.layers.Dropout(dropout_prob, name='last_dropout')(x)
    x = tf.keras.layers.ReLU(6.0, name='last_relu')(x)

    if pooling == 'avg':
        x = tf.keras.layers.GlobalAveragePooling2D()(x)
    elif pooling == 'max':
        x = tf.keras.layers.GlobalMaxPooling2D()(x)

    x = tf.keras.layers.Dense(num_classes,
                              activation='softmax',
                              use_bias=True,
                              name='logits')(x)

    return tf.keras.models.Model(inputs=img_input, outputs=x)
예제 #6
0
    def __init__(self,
                 n_classes,
                 num_layers=20,
                 num_initial_filters=16,
                 shortcut_connection=True,
                 weight_decay=2e-4,
                 batch_norm_momentum=0.99,
                 batch_norm_epsilon=1e-3,
                 pre_hidden_size = 256,
                 batch_norm_center=True,
                 batch_norm_scale=True,
                 group_norm_groups=16):
        """Constructor.
        Args:
          num_layers: int scalar, num of layers.
          shortcut_connection: bool scalar, whether to add shortcut connection in
            each Resnet unit. If False, degenerates to a 'Plain network'.
          weight_decay: float scalar, weight for l2 regularization.
          batch_norm_momentum: float scalar, the moving avearge decay.
          batch_norm_epsilon: float scalar, small value to avoid divide by zero.
          batch_norm_center: bool scalar, whether to center in the batch norm.
          batch_norm_scale: bool scalar, whether to scale in the batch norm.
        """
        super().__init__()
        if num_layers not in (20, 26, 32, 44, 56, 110):
            raise ValueError('num_layers must be one of 20, 32, 44, 56 or 110.')

        self._num_layers = num_layers
        self._num_initial_filters = num_initial_filters
        self._shortcut_connection = shortcut_connection
        self._weight_decay = weight_decay
        self._batch_norm_momentum = batch_norm_momentum
        self._batch_norm_epsilon = batch_norm_epsilon
        self._batch_norm_center = batch_norm_center
        self._batch_norm_scale = batch_norm_scale
        self._group_norm_groups = group_norm_groups

        self._num_units = (num_layers - 2) // 6

        self._kernel_regularizer = regularizers.l2(weight_decay)

        self._init_conv = layers.Conv2D(self._num_initial_filters, 3, 1, 'same', use_bias=False,
                                        kernel_regularizer=self._kernel_regularizer, name='init_conv')

        self._block1 = models.Sequential([ResNetUnit(
            self._num_initial_filters,
            1,
            shortcut_connection,
            True if i == 0 else False,
            weight_decay,
            batch_norm_momentum,
            batch_norm_epsilon,
            batch_norm_center,
            batch_norm_scale,
            group_norm_groups,
            'res_net_unit_%d' % (i + 1)) for i in range(self._num_units)],
            name='block1')
        self._block2 = models.Sequential([ResNetUnit(
            self._num_initial_filters*2,
            2 if i == 0 else 1,
            shortcut_connection,
            False if i == 0 else False,
            weight_decay,
            batch_norm_momentum,
            batch_norm_epsilon,
            batch_norm_center,
            batch_norm_scale,
            group_norm_groups,
            'res_net_unit_%d' % (i + 1)) for i in range(self._num_units)],
            name='block2')
        self._block3 = models.Sequential([ResNetUnit(
            self._num_initial_filters*4,
            2 if i == 0 else 1,
            shortcut_connection,
            False if i == 0 else False,
            weight_decay,
            batch_norm_momentum,
            batch_norm_epsilon,
            batch_norm_center,
            batch_norm_scale,
            group_norm_groups,
            'res_net_unit_%d' % (i + 1)) for i in range(self._num_units)],
            name='block3')

        self._global_avg = tf.keras.layers.GlobalAveragePooling2D(name="GlobalAvgPooling")

        self._last_bn = tfalayers.GroupNormalization(groups=group_norm_groups,
                                                     axis=-1,
                                                     epsilon=batch_norm_epsilon,
                                                     center=batch_norm_center,
                                                     scale=batch_norm_scale)
        self._pre_hidden = tf.keras.layers.Dense(pre_hidden_size, use_bias=False, activation='linear',
                                                 kernel_regularizer=self._kernel_regularizer)

        self._mlp_dense_out = tf.keras.layers.Dense(n_classes, use_bias=False, name="Head_Dense", activation='softmax',
                                                    kernel_regularizer=self._kernel_regularizer)