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