def mbconv_block(inputs, in_filters, out_filters, expand_ratio, strides, use_projection, kernel_size=3, se_ratio=None, batch_norm_relu=nn_ops.BatchNormRelu(), dropblock=nn_ops.Dropblock(), drop_connect_rate=None, data_format='channels_last', is_training=False): """The bottleneck block with BN and DropBlock after convolutions. Args: inputs: a `Tensor` of size `[batch, channels, height, width]`. in_filters: a `int` number of filters for the input feature map. out_filters: a `int` number of filters for the output feature map. expand_ratio: a `int` number as the feature dimension expansion ratio. strides: a `int` block stride. If greater than 1, this block will ultimately downsample the input. use_projection: a `bool` for whether this block should use a projection shortcut (versus the default identity shortcut). This is usually `True` for the first block of a block group, which may change the number of filters and the resolution. kernel_size: kernel size for the depthwise convolution. se_ratio: squeeze and excitation ratio. batch_norm_relu: an operation that is added after convolutions, including a batch norm layer and an optional relu activation. dropblock: a drop block layer that is added after convluations. Note that the default implementation does not apply any drop block. drop_connect_rate: a 'float' number that specifies the drop connection rate of the block. Note that the default `None` means no drop connection is applied. data_format: a `str` that specifies the data format. is_training: a `bool` if True, the model is in training mode. Returns: The output `Tensor` of the block. """ tf.logging.info('-----> Building mbconv block.') shortcut = inputs if use_projection: shortcut = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=out_filters, kernel_size=1, strides=strides, data_format=data_format) shortcut = batch_norm_relu(shortcut, is_training=is_training) shortcut = dropblock(shortcut, is_training=is_training) # First 1x1 conv for channel expansion. inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=in_filters * expand_ratio, kernel_size=1, strides=1, data_format=data_format) inputs = batch_norm_relu(inputs, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) # Second depthwise conv. inputs = nn_ops.depthwise_conv2d_fixed_padding(inputs=inputs, kernel_size=kernel_size, strides=strides, data_format=data_format) inputs = batch_norm_relu(inputs, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) # Squeeze and excitation. if se_ratio is not None and se_ratio > 0 and se_ratio <= 1: inputs = nn_ops.squeeze_excitation(inputs, in_filters, se_ratio, expand_ratio=expand_ratio, data_format=data_format) # Third 1x1 conv for reversed bottleneck. inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=out_filters, kernel_size=1, strides=1, data_format=data_format) inputs = batch_norm_relu(inputs, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) if drop_connect_rate: inputs = nn_ops.drop_connect(inputs, is_training, drop_connect_rate) return tf.add(inputs, shortcut)
def bottleneck_block(inputs, filters, strides, use_projection, activation=tf.nn.relu, batch_norm_relu=nn_ops.BatchNormRelu(), dropblock=nn_ops.Dropblock(), drop_connect_rate=None, data_format='channels_last', is_training=False): """The bottleneck block with BN and DropBlock after convolutions. Args: inputs: a `Tensor` of size `[batch, channels, height, width]`. filters: a `int` number of filters for the first two convolutions. Note that the third and final convolution will use 4 times as many filters. strides: an `int` block stride. If greater than 1, this block will ultimately downsample the input. use_projection: a `bool` for whether this block should use a projection shortcut (versus the default identity shortcut). This is usually `True` for the first block of a block group, which may change the number of filters and the resolution. activation: activation function. Support 'relu' and 'swish'. batch_norm_relu: an operation that is added after convolutions, including a batch norm layer and an optional relu activation. dropblock: a drop block layer that is added after convluations. Note that the default implementation does not apply any drop block. drop_connect_rate: a 'float' number that specifies the drop connection rate of the block. Note that the default `None` means no drop connection is applied. data_format: a `str` that specifies the data format. is_training: a `bool` if True, the model is in training mode. Returns: The output `Tensor` of the block. """ logging.info('-----> Building bottleneck block.') shortcut = inputs if use_projection: out_filters = 4 * filters shortcut = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=out_filters, kernel_size=1, strides=strides, data_format=data_format) shortcut = batch_norm_relu(shortcut, relu=False, is_training=is_training) shortcut = dropblock(shortcut, is_training=is_training) inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=filters, kernel_size=1, strides=1, data_format=data_format) inputs = batch_norm_relu(inputs, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=filters, kernel_size=3, strides=strides, data_format=data_format) inputs = batch_norm_relu(inputs, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=4 * filters, kernel_size=1, strides=1, data_format=data_format) inputs = batch_norm_relu(inputs, relu=False, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) if drop_connect_rate: inputs = nn_ops.drop_connect(inputs, is_training, drop_connect_rate) return activation(inputs + shortcut)
def fused_mbconv_block(inputs, in_filters, out_filters, expand_ratio, strides, kernel_size=3, se_ratio=None, batch_norm_activation=nn_ops.BatchNormActivation(), dropblock=nn_ops.Dropblock(), drop_connect_rate=None, data_format='channels_last', is_training=False): """The fused bottleneck block with BN and DropBlock after convolutions. Args: inputs: a `Tensor` of size `[batch, channels, height, width]`. in_filters: a `int` number of filters for the input feature map. out_filters: a `int` number of filters for the output feature map. expand_ratio: a `int` number as the feature dimension expansion ratio. strides: a `int` block stride. If greater than 1, this block will ultimately downsample the input. kernel_size: kernel size for the depthwise convolution. se_ratio: squeeze and excitation ratio. batch_norm_activation: an operation that includes a batch normalization layer followed by an optional activation layer. dropblock: a drop block layer that is added after convluations. Note that the default implementation does not apply any drop block. drop_connect_rate: a 'float' number that specifies the drop connection rate of the block. Note that the default `None` means no drop connection is applied. data_format: a `str` that specifies the data format. is_training: a `bool` if True, the model is in training mode. Returns: The output `Tensor` of the block. """ tf.logging.info('-----> Building fused mbconv block.') shortcut = inputs # First 1x1 conv for channel expansion. inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=in_filters * expand_ratio, kernel_size=kernel_size, strides=strides, data_format=data_format) inputs = batch_norm_activation(inputs, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) # Squeeze and excitation. if se_ratio is not None and se_ratio > 0 and se_ratio <= 1: inputs = nn_ops.squeeze_excitation(inputs, in_filters, se_ratio, expand_ratio=expand_ratio, data_format=data_format) # Third 1x1 conv for reversed bottleneck. inputs = nn_ops.conv2d_fixed_padding(inputs=inputs, filters=out_filters, kernel_size=1, strides=1, data_format=data_format) inputs = batch_norm_activation(inputs, relu=False, is_training=is_training) inputs = dropblock(inputs, is_training=is_training) if in_filters == out_filters and strides == 1: if drop_connect_rate: inputs = nn_ops.drop_connect(inputs, is_training, drop_connect_rate) inputs = tf.add(inputs, shortcut) return inputs