def _make_shared_convs(self):
self.box_shared_convs = tf.keras.Sequential(name="box_net")
self.class_shared_convs = tf.keras.Sequential(name="cls_net")
i = 0
for _ in range(self.cfg.repeats):
self.box_shared_convs.add(
tf.keras.layers.Conv2D(filters=self.cfg.feat_dims,
kernel_size=(3, 3),
padding="same",
strides=(1, 1),
use_bias=True,
name="%d" % i))
self.class_shared_convs.add(
tf.keras.layers.Conv2D(filters=self.cfg.feat_dims,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="%d" % i))
i += 1
self.box_shared_convs.add(
build_normalization(name="%d" % i,
**self.cfg.normalization.as_dict()))
self.class_shared_convs.add(
build_normalization(name="%d" % i,
**self.cfg.normalization.as_dict()))
i += 1
self.box_shared_convs.add(
build_activation(name="%d" % i,
**self.cfg.activation.as_dict()))
self.class_shared_convs.add(
build_activation(name="%d" % i,
**self.cfg.activation.as_dict()))
i += 1
def __init__(self,
convolution,
filters,
strides=1,
dilation_rate=1,
normalization="group_norm",
group=16,
activation="relu",
weight_decay=0.,
use_conv_shortcut=False,
**kwargs):
super(BasicBlock, self).__init__(**kwargs)
axis = 3 if tf.keras.backend.image_data_format(
) == "channels_last" else 1
self.conv1 = build_convolution(
convolution,
filters=filters,
kernel_size=3,
strides=strides,
padding="same",
dilation_rate=dilation_rate if strides == 1 else 1,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
use_bias=False)
self.norm1 = build_normalization(normalization, axis=axis, group=group)
self.act = tf.keras.layers.Activation(activation)
self.conv2 = build_convolution(
convolution,
filters=filters,
kernel_size=3,
strides=1,
padding="same",
dilation_rate=dilation_rate,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
use_bias=False)
self.norm2 = build_normalization(normalization, axis=axis, group=group)
if use_conv_shortcut:
if strides >= 2:
self.avg_pool = tf.keras.layers.AvgPool2D(2, strides, "same")
self.conv3 = build_convolution(
convolution,
filters=filters,
kernel_size=1,
strides=1,
padding="same",
kernel_regularizer=tf.keras.regularizers.l2(weight_decay))
self.norm3 = build_normalization(normalization,
axis=axis,
group=group)
self.use_conv_shortcut = use_conv_shortcut
def __init__(self,
convolution="conv2d",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="swish"),
feat_dims=64,
conv_bn_act_pattern=False,
data_format="channels_last",
name='op_after_combine'):
super(OpAfterCombine, self).__init__(name=name)
self.feat_dims = feat_dims
self.data_format = data_format
self.conv_bn_act_pattern = conv_bn_act_pattern
self.conv_op = build_convolution(convolution=convolution,
filters=feat_dims,
kernel_size=(3, 3),
padding='same',
use_bias=not self.conv_bn_act_pattern,
data_format=self.data_format,
name='conv')
self.bn = build_normalization(**normalization, name='bn')
self.act = build_activation(**activation)
def build(self, input_shape):
if self.padding == "same":
p = ((self.kernel_size[0] - 1) // 2,
(self.kernel_size[1] - 1) // 2)
self.pad = tf.keras.layers.ZeroPadding2D(
p, data_format=self.data_format)
self.conv = tf.keras.layers.DepthwiseConv2D(
kernel_size=self.kernel_size,
strides=self.strides,
padding="valid",
data_format=self.data_format,
dilation_rate=self.dilation_rate,
use_bias=self.use_bias,
trainable=self.trainable,
kernel_initializer=self.kernel_initializer,
name="conv2d")
self.norm_func = (build_normalization(
name=self.normalization["normalization"], **self.normalization)
if self.normalization is not None else None)
self.act_func = (build_activation(name=self.activation["activation"],
**self.activation)
if self.activation is not None else None)
self.dropblock_func = (DropBlock2D(name="dropblock", **self.dropblock)
if self.dropblock is not None else None)
def transition_block(self, x, reduction, strides, trainable, name):
"""A transition block.
Args:
x: input tensor.
reduction: float, compression rate at transition layers.
strides: integer, stride in pool layer.
trainable: bool, does freeze this block.
name: string, block label.
Returns:
output tensor for the block.
"""
bn_axis = 3 if tf.keras.backend.image_data_format(
) == "channels_last" else 1
x = build_normalization(**self.normalization, name=name + "_bn")(x)
x = tf.keras.layers.Activation(**self.activation,
name=name + "_relu")(x)
preact = x
x = build_convolution(
self.convolution,
filters=int(tf.keras.backend.int_shape(x)[bn_axis] * reduction),
kernel_size=1,
use_bias=False,
trainable=trainable,
name=name + "_conv")(x)
x = tf.keras.layers.AvgPool2D(2, strides=strides,
name=name + "_pool")(x)
return preact, x
def __init__(self,
feat_level,
target_num_channels,
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="swish"),
apply_bn=False,
conv_after_downsample=False,
data_format=None,
pooling_type=None,
upsampling_type=None,
name='resample_p0'):
super(ResampleFeatureMap, self).__init__(name=name)
self.apply_bn = apply_bn
self.data_format = data_format
self.target_num_channels = target_num_channels
self.feat_level = feat_level
self.conv_after_downsample = conv_after_downsample
self.pooling_type = pooling_type or 'max'
self.upsampling_type = upsampling_type or 'nearest'
self.conv2d = tf.keras.layers.Conv2D(self.target_num_channels, (1, 1),
padding='same',
data_format=self.data_format,
name='conv2d')
self.bn = build_normalization(**normalization, name='bn')
def resnet(self):
trainable = 0 not in self.frozen_stages
inputs = tf.keras.layers.Lambda(
function=lambda inp: inp - [123.68, 116.78, 103.94],
name="mean_subtraction")(self.img_input)
x = tf.keras.layers.ZeroPadding2D(((3, 3), (3, 3)),
name="conv1_pad")(inputs)
x = build_convolution(self.convolution,
filters=64,
kernel_size=(7, 7),
strides=self.strides[0],
padding="valid",
dilation_rate=self.dilation_rates[0],
use_bias=True,
trainable=trainable,
kernel_initializer="he_normal",
kernel_regularizer=tf.keras.regularizers.l2(
self.weight_decay),
name="conv1_conv")(x)
x = build_normalization(**self.normalization, name="conv1_bn")(x)
x1 = build_activation(**self.activation, name="conv1_relu")(x)
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)),
name="pool1_pad")(x1)
x = tf.keras.layers.MaxPool2D((3, 3),
self.strides[1],
"valid",
name="pool1_pool")(x)
trainable = 2 not in self.frozen_stages
x2 = self.stack(x, 64, 1, self.dilation_rates[1], trainable,
self.blocks[0], "conv2")
trainable = 3 not in self.frozen_stages
x3 = self.stack(x2, 128, self.strides[2], self.dilation_rates[2],
trainable, self.blocks[1], "conv3")
trainable = 4 not in self.frozen_stages
x4 = self.stack(x3, 256, self.strides[3], self.dilation_rates[3],
trainable, self.blocks[2], "conv4")
trainable = 5 not in self.frozen_stages
x5 = self.stack(x4, 512, self.strides[4], self.dilation_rates[4],
trainable, self.blocks[3], "conv5")
if self._is_classifier:
x = tf.keras.layers.GlobalAvgPool2D(name="avg_pool")(x5)
x = tf.keras.layers.Dropout(rate=self.drop_rate)(x)
outputs = tf.keras.layers.Dense(self.num_classes,
activation="softmax",
name="probs")(x)
return tf.keras.Model(inputs=self.img_input,
outputs=outputs,
name=self.name)
else:
outputs = [
o for i, o in enumerate([x1, x2, x3, x4, x5])
if i + 1 in self.output_indices
]
return outputs
def squeeze_excitation(inputs,
in_filters,
se_ratio=0.25,
data_format="channels_last",
normalization=None,
trainable=True,
name="se"):
"""Squeeze and excitation implementation.
Args:
inputs: `Tensor` of size `[batch, channels, height_in, width_in]`.
in_filters: `int` number of input filteres before expansion.
se_ratio: `float` a se ratio between 0 and 1 for squeeze and excitation.
data_format: An optional string from: "channels_last", "channels_first".
Defaults to "channels_last".
Returns:
A `Tensor` of shape `[batch, filters, height_out, width_out]`.
"""
if isinstance(se_ratio, float):
num_reduced_filters = max(1, int(in_filters * se_ratio))
elif isinstance(se_ratio, int):
num_reduced_filters = max(1, in_filters // se_ratio)
else:
raise ValueError("se_ratio should be `float` or `int`")
# Process input
if data_format == "channels_first":
spatial_dims = [2, 3]
else:
spatial_dims = [1, 2]
x = tf.keras.layers.Lambda(
lambda inp: tf.reduce_mean(inp, spatial_dims, keepdims=True),
name=name + "/global_avgpool")(inputs)
x = tf.keras.layers.Conv2D(
num_reduced_filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=tf.keras.initializers.VarianceScaling(),
padding="same",
activation="relu" if normalization is None else None,
data_format=data_format,
trainable=trainable,
name=name + "/squeeze")(x)
if normalization is not None:
x = build_normalization(**normalization,
name=name + "/squeeze/norm")(x)
x = tf.keras.layers.ReLU(name=name + "/squeeze/relu")(x)
x = tf.keras.layers.Conv2D(
in_filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=tf.keras.initializers.VarianceScaling(),
padding="same",
data_format=data_format,
activation="sigmoid",
trainable=trainable,
use_bias=True,
name=name + "/excitation")(x)
return tf.keras.layers.Multiply(name=name + "/multiply")([x, inputs])
def _make_shared_convs(self):
self.box_shared_convs = [
build_convolution(convolution=self.cfg.convolution,
filters=self.cfg.feat_dims,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=self.cfg.normalization is None
or self.cfg.convolution == "separable_conv2d",
name="box_net/box-%d" % i)
for i in range(self.cfg.repeats)
]
if self.cfg.normalization:
self.box_norm_layers = {
"level%d" % level: [
build_normalization(**self.cfg.normalization.as_dict(),
name="box_net/box-%d-bn-%d" %
(i, level))
for i in range(self.cfg.repeats)
]
for level in range(self.min_level, self.max_level + 1)
}
self.class_shared_convs = [
build_convolution(convolution=self.cfg.convolution,
filters=self.cfg.feat_dims,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=self.cfg.normalization is None
or self.cfg.convolution == "separable_conv2d",
name="class_net/class-%d" % i)
for i in range(self.cfg.repeats)
]
if self.cfg.normalization:
self.class_norm_layers = {
"level%d" % level: [
build_normalization(**self.cfg.normalization.as_dict(),
name="class_net/class-%d-bn-%d" %
(i, level))
for i in range(self.cfg.repeats)
]
for level in range(self.min_level, self.max_level + 1)
}
self.act_fn = build_activation(**self.cfg.activation.as_dict())
def conv_block(self, x, growth_rate, dilation_rate, drop_rate, trainable,
name):
"""A building block for a dense block.
Args:
x: input tensor.
growth_rate: float, growth rate at dense layers.
dilation_rate: integer, dilation rate.
drop_rate: float, the dropout rate.
trainable: bool, does freeze this block
name: string, block label.
Returns:
Output tensor for the block.
"""
bn_axis = 3 if tf.keras.backend.image_data_format(
) == "channels_last" else 1
x1 = build_normalization(**self.normalization, name=name + "_0_bn")(x)
x1 = tf.keras.layers.Activation(**self.activation,
name=name + "_0_relu")(x1)
x1 = build_convolution(self.convolution,
filters=4 * growth_rate,
kernel_size=1,
use_bias=False,
trainable=trainable,
name=name + "_1_conv")(x1)
x1 = build_normalization(**self.normalization, name=name + "_1_bn")(x1)
x1 = tf.keras.layers.Activation(**self.activation,
name=name + "_1_relu")(x1)
x1 = build_convolution(self.convolution,
filters=growth_rate,
kernel_size=3,
padding="same",
dilation_rate=dilation_rate,
use_bias=False,
trainable=trainable,
name=name + "_2_conv")(x1)
x1 = tf.keras.layers.Dropout(rate=drop_rate)(x1)
x = tf.keras.layers.Concatenate(axis=bn_axis,
name=name + "_concat")([x, x1])
return x
def build(self, input_shape):
if self.strides != (1, 1) and self.padding == "same":
p = ((self.kernel_size[0] - 1) // 2,
(self.kernel_size[1] - 1) // 2)
self.pad = tf.keras.layers.ZeroPadding2D(
p, data_format=self.data_format)
self.padding = "valid"
self.conv = tf.keras.layers.Conv2D(
filters=self.filters,
kernel_size=self.kernel_size,
strides=self.strides,
padding=self.padding,
data_format=self.data_format,
dilation_rate=self.dilation_rate,
groups=self.groups,
use_bias=self.use_bias,
trainable=self.trainable,
kernel_initializer=self.kernel_initializer,
name="conv2d")
if self.normalization is not None:
if self.gamma_zeros:
self.norm_func = build_normalization(
name=self.normalization["normalization"],
**self.normalization)
else:
self.norm_func = build_normalization(
gamma_initializer="zeros",
name=self.normalization["normalization"],
**self.normalization)
else:
self.norm_func = None
self.act_func = (build_activation(name=self.activation["activation"],
**self.activation)
if self.activation is not None else None)
self.dropblock_func = (DropBlock2D(**self.dropblock, name="dropblock")
if self.dropblock is not None else None)
self.built = True
def prediction_head(inputs,
shared_convolutions,
normalization,
activation="swish",
repeats=3,
level=3,
name="prediction_head"):
x = inputs
for i in range(repeats):
x = shared_convolutions[i](x)
if normalization is not None:
x = build_normalization(normalization=normalization.normalization,
name=name + "-%d-bn-%d" % (i, level))(x)
if activation is not None:
x = build_activation(activation=activation.activation,
name=name + "-%d-%s-%d" %
(i, activation.activation, level))(x)
return x
def dcn_deconv(inputs,
filters,
deconv_kernel_size,
deconv_strides=2,
deconv_padding="same",
deconv_output_padding=None,
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
kernel_initializer="he_normal",
groups=1,
dilation_rate=1,
name="deconv"):
x = ConvNormActBlock(filters=filters,
kernel_size=3,
dilation_rate=dilation_rate,
normalization=normalization,
activation=activation,
kernel_initializer=kernel_initializer,
use_bias=True,
name=name + "/conv")(inputs)
deconv_kernel_shape = [
deconv_kernel_size, deconv_kernel_size, filters, filters
]
x = tf.keras.layers.Conv2DTranspose(
filters=filters,
kernel_size=deconv_kernel_size,
strides=deconv_strides,
padding=deconv_padding,
output_padding=deconv_output_padding,
use_bias=False,
kernel_initializer=tf.keras.initializers.Constant(
_deconv_init(deconv_kernel_shape)),
name=name + "/up_sample")(x)
x = build_normalization(name=name + "/up_bn", **normalization)(x)
x = build_activation(name=name + "/relu", **activation)(x)
return x
def build_model(self):
# Build stem
x = self.img_input
x = layers.Rescaling(1. / 255.)(x)
x = layers.Normalization(axis=self.bn_axis)(x)
x = tf.keras.layers.ZeroPadding2D(padding=imagenet_utils.correct_pad(
x, 3),
name='stem_conv_pad')(x)
x = tf.keras.layers.Conv2D(self.round_filters(32),
3,
strides=2,
padding='valid',
use_bias=False,
kernel_initializer=CONV_KERNEL_INITIALIZER,
name='stem_conv')(x)
x = build_normalization(**self.normalization, name='stem_bn')(x)
x = build_activation(**self.activation, name='stem_activation')(x)
block_outputs = [x]
# Build blocks
blocks_args = copy.deepcopy(self.blocks_args)
b = 0
blocks = float(sum(args['repeats'] for args in blocks_args))
output_stages = [1, 2, 4, 6]
for (i, args) in enumerate(blocks_args):
assert args['repeats'] > 0
# Update block input and output filters based on depth multiplier.
args['filters_in'] = self.round_filters(args['filters_in'])
args['filters_out'] = self.round_filters(args['filters_out'])
strides = args["strides"]
for j in range(self.round_repeats(args.pop('repeats'))):
# The first block needs to take care of stride and filter size increase.
if j > 0:
args['strides'] = 1
args['filters_in'] = args['filters_out']
x = self.block(x,
self.activation,
self.drop_connect_rate * b / blocks,
name='block{}{}_'.format(i + 1, chr(j + 97)),
**args)
b += 1
if i in output_stages:
block_outputs.append(x)
# Build top
x = tf.keras.layers.Conv2D(self.round_filters(1280),
1,
padding='same',
use_bias=False,
kernel_initializer=CONV_KERNEL_INITIALIZER,
name='top_conv')(x)
x = build_normalization(**self.normalization, name='top_bn')(x)
x = build_activation(**self.activation, name='top_activation')(x)
if -1 in self.output_indices:
x = tf.keras.layers.GlobalAveragePooling2D(name='avg_pool')(x)
if self.dropout_rate > 0:
x = tf.keras.layers.Dropout(self.dropout_rate,
name='top_dropout')(x)
x = tf.keras.layers.Dense(
1000,
activation=self.classifier_activation,
kernel_initializer=DENSE_KERNEL_INITIALIZER,
name='predictions')(x)
# Ensure that the model takes into account
# any potential predecessors of `input_tensor`.
if self.input_tensor is not None:
inputs = tf.keras.utils.get_source_inputs(self.input_tensor)
else:
inputs = self.img_input
# Create model.
return tf.keras.Model(inputs=self.img_input,
outputs=x,
name=model_name)
return [block_outputs[i - 1] for i in self.output_indices]
def build_model(self):
# rgb_mean = self._rgb_mean
# rgb_std = self._rgb_std
def _norm(inp):
rgb_mean = tf.constant([0.485 * 255, 0.456 * 255, 0.406 * 255],
inp.dtype, [1, 1, 1, 3])
rgb_std = tf.constant([0.229 * 255, 0.224 * 255, 0.225 * 255],
inp.dtype, [1, 1, 1, 3])
return (inp - rgb_mean) * (1. / rgb_std)
x = tf.keras.layers.Lambda(function=_norm,
name="norm_input")(self.img_input)
if self.use_resnetd_stem:
x = ConvNormActBlock(filters=32,
kernel_size=3,
strides=self.strides[0],
dilation_rate=self.dilation_rates[0]
if self.strides[0] == 1 else 1,
trainable=1 not in self.frozen_stages,
kernel_initializer=self.kernel_initializer,
normalization=self.normalization,
activation=self.activation,
data_format=self.data_format,
name="stem/conv1")(x)
x = ConvNormActBlock(filters=32,
kernel_size=3,
strides=1,
dilation_rate=self.dilation_rates[0],
trainable=1 not in self.frozen_stages,
kernel_initializer=self.kernel_initializer,
normalization=self.normalization,
activation=self.activation,
data_format=self.data_format,
name="stem/conv2")(x)
x = ConvNormActBlock(
filters=64,
kernel_size=3,
strides=1,
dilation_rate=self.dilation_rates[0],
trainable=1 not in self.frozen_stages,
kernel_initializer=self.kernel_initializer,
normalization=None
if self.preactivation else self.normalization,
activation=None if self.preactivation else self.activation,
data_format=self.data_format,
name="stem/conv3")(x)
else:
x = ConvNormActBlock(
filters=64,
kernel_size=7,
strides=self.strides[0],
dilation_rate=self.dilation_rates[0]
if self.strides[0] == 1 else 1,
trainable=1 not in self.frozen_stages,
kernel_initializer=self.kernel_initializer,
normalization=None
if self.preactivation else self.normalization,
activation=None if self.preactivation else self.activation,
data_format=self.data_format,
name="stem/conv1")(x)
if not self.skip_stem_max_pool:
if self.replace_stem_max_pool:
x = ConvNormActBlock(
filters=64,
kernel_size=3,
strides=self.strides[1],
dilation_rate=1,
trainable=1 not in self.frozen_stages,
kernel_initializer=self.kernel_initializer,
normalization=None
if self.preactivation else self.normalization,
activation=None if self.preactivation else self.activation,
data_format=self.data_format,
name="stem/maxpool")(x)
else:
x = tf.keras.layers.MaxPool2D(3,
self.strides[1],
"same",
self.data_format,
name="stem/maxpool")(x)
self.in_filters = 64
trainable = 2 not in self.frozen_stages
x = self.stack(x, 64,
self.strides[1] if self.skip_stem_max_pool else 1,
self.dilation_rates[1], trainable, self.blocks[0],
"layer1")
trainable = 3 not in self.frozen_stages
x = self.stack(x, 128, self.strides[2], self.dilation_rates[2],
trainable, self.blocks[1], "layer2")
trainable = 4 not in self.frozen_stages
x = self.stack(x, 256, self.strides[3], self.dilation_rates[3],
trainable, self.blocks[2], "layer3")
trainable = 5 not in self.frozen_stages
x = self.stack(x, 512, self.strides[4], self.dilation_rates[4],
trainable, self.blocks[3], "layer4")
if self.preactivation:
x = build_normalization(**self.normalization)(x)
x = build_activation(**self.activation)(x)
x = tf.keras.layers.GlobalAvgPool2D(data_format=self.data_format)(x)
x = tf.keras.layers.Dropout(rate=self.drop_rate)(x)
x = tf.keras.layers.Dense(self.feat_dim, name="logits")(x)
model = tf.keras.Model(inputs=self.img_input,
outputs=x,
name=self.name)
return model
def basic_block(inputs,
filters,
strides=1,
dilation_rate=1,
data_format="channels_last",
kernel_initializer=tf.keras.initializers.VarianceScaling(),
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
init_gamma_zero=True,
activation=dict(activation="relu"),
trainable=True,
dropblock=None,
use_conv_shortcut=False,
expansion=1,
drop_connect_rate=None,
se_ratio=None,
preactivation=False,
use_resnetd_shortcut=True,
name=None,
**kwargs):
"""
Basic Residual block
Args:
filters(int): integer, filters of the bottleneck layer.
strides(int): default 1, stride of the first layer.
dilation_rate(int): default 1, dilation rate in 3x3 convolution.
data_format(str): default channels_last,
normalization(dict): the normalization name and hyper-parameters, e.g.
dict(normalization="batch_norm", momentum=0.9, epsilon=1e-3, axis=-1, trainable=True),
dict(normalization="group_norm", epsilon=1e-3, axis=-1) etc.
activation: the activation layer name.
trainable: does this block is trainable.
dropblock: the arguments in DropBlock2D
use_conv_shortcut: default True, use convolution shortcut if True,
otherwise identity shortcut.
name: string, block label, default None.
"""
if preactivation:
x = build_normalization(name=name + "/norm1", **normalization)(inputs)
x = build_activation(name=name + "/%s" % activation["activation"],
**activation)(x)
else:
x = inputs
shortcut = x
x = ConvNormActBlock(filters=filters,
kernel_size=3,
strides=strides,
data_format=data_format,
dilation_rate=1 if strides > 1 else dilation_rate,
kernel_initializer=kernel_initializer,
trainable=trainable,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/conv1")(x)
x = ConvNormActBlock(
filters=filters,
kernel_size=3,
strides=1,
data_format=data_format,
dilation_rate=dilation_rate,
kernel_initializer=kernel_initializer,
trainable=trainable,
normalization=None if preactivation else normalization,
activation=None,
dropblock=dropblock,
name=name + "/conv2")(x)
if use_conv_shortcut:
if use_resnetd_shortcut and strides > 1:
shortcut = tf.keras.layers.AvgPool2D(pool_size=strides,
strides=strides,
padding="same",
data_format=data_format,
name=name +
"/avgpool")(shortcut)
strides = 1
shortcut = ConvNormActBlock(
filters=filters,
kernel_size=1,
strides=strides,
data_format=data_format,
trainable=trainable,
kernel_initializer=kernel_initializer,
normalization=None if preactivation else normalization,
activation=None,
dropblock=dropblock,
gamma_zeros=init_gamma_zero,
name=name + "/shortcut")(shortcut)
if preactivation:
return tf.keras.layers.Add(name=name + "/add")([x, shortcut])
if se_ratio is not None and se_ratio > 0 and se_ratio <= 1:
x = squeeze_excitation(x,
in_filters=filters * expansion,
se_ratio=se_ratio,
data_format="channels_last",
name=name + "/se")
if drop_connect_rate:
x = DropConnect(drop_connect_rate, name=name + "/drop_connect")(x)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
x = build_activation(**activation,
name=name + "/" + activation["activation"])(x)
return x
def mbconv_block(inputs,
global_params,
block_args,
normalization,
drop_connect_rate=None,
trainable=True,
l2_regularizer=tf.keras.regularizers.l2(l=4e-5),
name=""):
expand_ratio = block_args.expand_ratio
data_format = global_params.data_format
_momentum = global_params.batch_norm_momentum
_epsilon = global_params.batch_norm_epsilon
_axis = normalization["axis"]
_mean_axis = [1, 2] if _axis == -1 or _axis == 3 else [2, 3]
filters = block_args.in_filters * expand_ratio
expand_ratio = expand_ratio
if expand_ratio != 1:
x = tf.keras.layers.Conv2D(
filters=filters,
kernel_size=(1, 1),
strides=(1, 1),
padding="same",
data_format=data_format,
use_bias=False,
name=name + "/conv2d",
trainable=trainable,
kernel_initializer=conv2d_kernel_initializer)(inputs)
x = build_normalization(**normalization,
name=name + "/batch_normalization")(x)
x = tf.keras.layers.Activation("swish", name=name + "/swish")(x)
else:
x = inputs
# Depthwise Convolution
# if block_args.strides == 2:
# x = tf.keras.layers.ZeroPadding2D(
# padding=imagenet_utils.correct_pad(x, block_args.kernel_size),
# name=name + '/dwconv_pad')(x)
# conv_pad = 'valid'
# else:
# conv_pad = 'same'
x = tf.keras.layers.DepthwiseConv2D(
kernel_size=block_args.kernel_size,
strides=block_args.strides,
padding="same",
data_format=data_format,
use_bias=False,
trainable=trainable,
depthwise_initializer=conv2d_kernel_initializer,
name=name + "/depthwise_conv2d")(x)
x = build_normalization(**normalization,
name=name + "/batch_normalization" if expand_ratio
== 1 else name + "/batch_normalization_1")(x)
x = tf.keras.layers.Activation(
"swish",
name=name + "/swish" if expand_ratio == 1 else name + "/swish_1")(x)
has_se = block_args.se_ratio is not None and 0 < block_args.se_ratio < 1
if has_se:
squeeze_filters = max(1,
int(block_args.in_filters * block_args.se_ratio))
se = tf.keras.layers.Lambda(
lambda inp: tf.reduce_mean(inp, axis=_mean_axis, keepdims=True),
name=name + "/se/global_pooling")(x)
se = tf.keras.layers.Conv2D(
filters=squeeze_filters,
kernel_size=(1, 1),
strides=(1, 1),
padding="same",
data_format=data_format,
use_bias=True,
kernel_initializer=conv2d_kernel_initializer,
trainable=trainable,
name=name + "/se/conv2d")(se)
se = tf.keras.layers.Activation("swish", name=name + "/se/swish_1")(se)
se = tf.keras.layers.Conv2D(
filters=filters,
kernel_size=(1, 1),
strides=(1, 1),
padding="same",
data_format=data_format,
use_bias=True,
trainable=trainable,
kernel_initializer=conv2d_kernel_initializer,
name=name + "/se/conv2d_1")(se)
se = tf.keras.layers.Activation("sigmoid",
name=name + "/se/sigmoid")(se)
x = tf.keras.layers.Multiply(name=name + "/se/multiply")([se, x])
x = tf.keras.layers.Conv2D(
block_args.out_filters,
kernel_size=(1, 1),
strides=(1, 1),
padding="same",
data_format=data_format,
use_bias=False,
trainable=trainable,
kernel_initializer=conv2d_kernel_initializer,
name=name + "/conv2d" if expand_ratio == 1 else name + "/conv2d_1")(x)
x = build_normalization(**normalization,
name=name + "/batch_normalization_2"
if expand_ratio > 1 else name +
"/batch_normalization_1")(x)
if block_args.id_skip:
if all(s == 1 for s in block_args.strides
) and block_args.in_filters == block_args.out_filters:
# x = DropConnect(drop_connect_rate, name=name + "/drop_connect")(x)
# x = tf.keras.layers.Dropout(drop_connect_rate, noise_shape=(None, 1, 1, 1), name=name + '/drop')(x)
x = tf.keras.layers.Add(name=name + "/add")([x, inputs])
return x
def build_model(self):
x = tf.keras.layers.Lambda(
lambda inp:
(inp - tf.constant([0.485, 0.456, 0.406], dtype=inp.dtype)) *
(1. / tf.constant([0.229, 0.224, 0.225], dtype=inp.dtype)),
name="mean_subtraction")(self.img_input)
x = tf.keras.layers.Conv2D(
round_filters(32, self.global_params),
kernel_size=(3, 3),
strides=(2, 2),
padding="same",
data_format=self.data_format,
use_bias=False,
kernel_initializer=conv2d_kernel_initializer,
trainable=1 not in self.frozen_stages,
kernel_regularizer=self.l2_regularizer,
name=self.name + "/stem/conv2d")(x)
x = build_normalization(**self.normalization,
name=self.name +
"/stem/batch_normalization")(x)
x = tf.keras.layers.Activation("swish",
name=self.name + "/stem/swish")(x)
block_outputs = []
for idx, b_args in enumerate(self.blocks):
drop_rate = self._drop_connect_rate
is_reduction = False
if b_args.super_pixel == 1 and idx == 0:
block_outputs.append(x)
elif (idx == self.num_blocks -
1) or self.blocks[idx + 1].strides[0] > 1:
is_reduction = True
if drop_rate:
drop_rate = 1.0 - drop_rate * float(idx) / self.num_blocks
x = mbconv_block(x,
global_params=self.global_params,
block_args=b_args,
normalization=self.normalization,
drop_connect_rate=drop_rate,
trainable=b_args.trainable,
l2_regularizer=self.l2_regularizer,
name=self.name + "/blocks_%d" % idx)
if is_reduction:
block_outputs.append(x)
if -1 in self.output_indices:
# Head part.
x = tf.keras.layers.Conv2D(
filters=round_filters(1280, self.global_params),
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=conv2d_kernel_initializer,
padding="same",
use_bias=False,
data_format=self.data_format,
kernel_regularizer=self.l2_regularizer,
name=self.name + "/head/conv2d")(x)
x = build_normalization(**self.normalization,
name=self.name +
"/head/batch_normalization")(x)
x = tf.keras.layers.Activation("swish",
name=self.name + "/head/swish")(x)
x = tf.keras.layers.GlobalAveragePooling2D(
data_format=self.data_format,
name=self.name + "/head/global_avg_pooling")(x)
x = tf.keras.layers.Dropout(self.global_params.dropout_rate,
name=self.name + "/head/dropout")(x)
x = tf.keras.layers.Dense(
self.global_params.num_classes,
activation="sigmoid",
kernel_initializer=dense_kernel_initializer,
name=self.name + "/head/dense")(x)
outputs = x
return tf.keras.Model(inputs=self.img_input,
outputs=outputs,
name=self.backbone_name)
else:
outputs = [block_outputs[i - 1] for i in self.output_indices]
return outputs
def resnet_v2(self):
trainable = 0 not in self.frozen_stages
x = tf.keras.layers.Lambda(function=lambda inp: inp / 127.5 - 1.,
name="mean_subtraction")(self.img_input)
x = tf.keras.layers.ZeroPadding2D(padding=((3, 3), (3, 3)),
name="conv1_pad")(x)
x1 = build_convolution(self.convolution,
filters=64,
kernel_size=7,
strides=self.strides[0],
padding="valid",
dilation_rate=self.dilation_rates[0],
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(
self.weight_decay),
use_bias=True,
name="conv1_conv")(x)
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)),
name="pool1_pad")(x1)
x = tf.keras.layers.MaxPool2D((3, 3),
strides=self.strides[1],
padding="valid",
name="pool1_pool")(x)
trainable = 1 not in self.frozen_stages
x, _ = self.stack(x, 64, self.strides[2], self.dilation_rates[1],
trainable, self.blocks[0], "conv2")
trainable = 2 not in self.frozen_stages
x, preact2 = self.stack(x, 128, self.strides[3],
self.dilation_rates[2], trainable,
self.blocks[1], "conv3")
trainable = 3 not in self.frozen_stages
x, preact3 = self.stack(x, 256, self.strides[4],
self.dilation_rates[3], trainable,
self.blocks[2], "conv4")
trainable = 4 not in self.frozen_stages
x, preact4 = self.stack(x, 512, 1, self.dilation_rates[4], trainable,
self.blocks[3], "conv5")
x = build_normalization(**self.normalization, name="post_bn")(x)
x5 = tf.keras.layers.Activation(**self.activation, name="post_relu")(x)
if self._is_classifier:
x = tf.keras.layers.GlobalAvgPool2D(name="avg_pool")(x5)
x = tf.keras.layers.Dropout(rate=self.drop_rate)(x)
x = tf.keras.layers.Dense(self.num_classes,
activation="softmax",
name="probs")(x)
return tf.keras.Model(inputs=self.img_input,
outputs=x,
name="resnet" + str(self.depth) + "v2")
outputs = [
o for i, o in enumerate([x1, preact2, preact3, preact4, x5])
if i in self.output_indices
]
return outputs
def build_model(self):
trainable = 0 not in self.frozen_stages
def _norm(inp):
inp = tf.image.convert_image_dtype(inp, tf.float32)
inp = inp / 127.5 - 1.
return inp
x = tf.keras.layers.Lambda(function=_norm,
name="norm_input")(self.img_input)
x = tf.keras.layers.ZeroPadding2D(padding=((3, 3), (3, 3)),
name="conv1_pad")(x)
x1 = build_convolution(self.convolution,
filters=64,
kernel_size=7,
strides=self.strides[0],
padding="valid",
dilation_rate=self.dilation_rates[0],
trainable=trainable,
use_bias=True,
name="conv1_conv")(x)
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)),
name="pool1_pad")(x1)
x = tf.keras.layers.MaxPool2D((3, 3),
strides=self.strides[1],
padding="valid",
name="pool1_pool")(x)
trainable = 1 not in self.frozen_stages
x, _ = self.stack(x, 64, self.strides[2], self.dilation_rates[1],
trainable, self.blocks[0], "conv2")
trainable = 2 not in self.frozen_stages
x, preact2 = self.stack(x, 128, self.strides[3],
self.dilation_rates[2], trainable,
self.blocks[1], "conv3")
trainable = 3 not in self.frozen_stages
x, preact3 = self.stack(x, 256, self.strides[4],
self.dilation_rates[3], trainable,
self.blocks[2], "conv4")
trainable = 4 not in self.frozen_stages
x, preact4 = self.stack(x, 512, 1, self.dilation_rates[4], trainable,
self.blocks[3], "conv5")
x = build_normalization(**self.normalization, name="post_bn")(x)
x5 = tf.keras.layers.Activation(**self.activation, name="post_relu")(x)
if -1 in self.output_indices:
x = tf.keras.layers.GlobalAvgPool2D(name="avg_pool")(x5)
x = tf.keras.layers.Dropout(rate=self.drop_rate)(x)
outputs = tf.keras.layers.Dense(self.num_classes, name="probs")(x)
else:
outputs = [
o for i, o in enumerate([x1, preact2, preact3, preact4, x5])
if i in self.output_indices
]
return tf.keras.Model(inputs=self.img_input,
outputs=outputs,
name=self.name)
def __init__(self,
filters,
kernel_size,
strides=1,
groups=1,
modulation=True,
padding="same",
data_format="channels_last",
dilation_rate=1,
kernel_initializer="he_normal",
trainable=True,
use_bias=False,
normalization=dict(normalization="batch_norm",
axis=-1,
trainable=True),
gamma_zeros=False,
activation=None,
dropblock=None,
name=None):
super(DCNNormActBlock, self).__init__(name=name)
assert isinstance(kernel_size, (int, tuple, list))
self.strides = strides
if isinstance(kernel_size, int):
kernel_size = (kernel_size, kernel_size)
if isinstance(strides, int):
strides = (strides, strides)
# if strides != (1, 1):
# dilation_rate = (1, 1)
if isinstance(dilation_rate, int):
dilation_rate = (dilation_rate, dilation_rate)
if strides == (1, 1):
p = 1
padding == "same"
else:
p = ((kernel_size[0] - 1) // 2 * dilation_rate[0],
(kernel_size[1] - 1) * dilation_rate[1] // 2)
self.pad = tf.keras.layers.ZeroPadding2D(p,
data_format=data_format)
padding = "valid"
self.conv = DCNv2(filters=filters,
kernel_size=kernel_size,
strides=strides,
padding=padding,
data_format=data_format,
dilation_rate=dilation_rate,
groups=groups,
use_bias=normalization is None or use_bias,
trainable=trainable,
kernel_initializer=kernel_initializer,
name="dcn")
self._nk = kernel_size * kernel_size if isinstance(
kernel_size, int) else kernel_size[0] * kernel_size[1]
if modulation:
offset_filters = self._nk * 3
self.offset_conv = tf.keras.layers.Conv2D(
filters=offset_filters,
kernel_size=(3, 3),
strides=strides,
padding=padding,
data_format=data_format,
dilation_rate=dilation_rate,
groups=groups,
use_bias=normalization is None or use_bias,
trainable=trainable,
kernel_initializer=kernel_initializer,
name="offset_conv2d")
else:
offset_filters = kernel_size * kernel_size * 2 if isinstance(
kernel_size, int) else kernel_size[0] * kernel_size[1] * 2
self.offset_conv = tf.keras.layers.Conv2D(
filters=filters,
kernel_size=(3, 3),
strides=strides,
padding=padding,
data_format=data_format,
dilation_rate=dilation_rate,
groups=groups,
use_bias=normalization is None or use_bias,
trainable=trainable,
kernel_initializer=kernel_initializer,
name="offset_conv2d")
self.modulation = modulation
channel_axis = -1 if data_format == "channels_last" else 1
if not trainable and normalization is not None:
normalization["trainable"] = False
if normalization is not None:
normalization["axis"] = channel_axis
self.norm = build_normalization(
**normalization,
gamma_initializer="zeros" if gamma_zeros else "ones",
name=normalization["normalization"])
else:
self.norm = None
# self.act = build_activation(**activation, name=activation["activation"]) if activation is not None else None
# self.dropblock = DropBlock2D(**dropblock, name="dropblock") if dropblock is not None else None
if activation is not None:
self.act = (build_activation(**activation) if isinstance(
activation, dict) else build_activation(activation=activation))
else:
self.act = None
def block(self,
inputs,
activation='swish',
drop_rate=0.,
name='',
filters_in=32,
filters_out=16,
kernel_size=3,
strides=1,
expand_ratio=1,
se_ratio=0.,
id_skip=True):
"""An inverted residual block.
Arguments:
inputs: input tensor.
activation: activation function.
drop_rate: float between 0 and 1, fraction of the input units to drop.
name: string, block label.
filters_in: integer, the number of input filters.
filters_out: integer, the number of output filters.
kernel_size: integer, the dimension of the convolution window.
strides: integer, the stride of the convolution.
expand_ratio: integer, scaling coefficient for the input filters.
se_ratio: float between 0 and 1, fraction to squeeze the input filters.
id_skip: boolean.
Returns:
output tensor for the block.
"""
# Expansion phase
filters = filters_in * expand_ratio
if expand_ratio != 1:
x = tf.keras.layers.Conv2D(
filters,
1,
padding='same',
use_bias=False,
kernel_initializer=CONV_KERNEL_INITIALIZER,
name=name + 'expand_conv')(inputs)
x = build_normalization(**self.normalization,
name=name + 'expand_bn')(x)
x = build_activation(**self.activation,
name=name + 'expand_activation')(x)
else:
x = inputs
# Depthwise Convolution
if strides == 2:
x = tf.keras.layers.ZeroPadding2D(
padding=imagenet_utils.correct_pad(x, kernel_size),
name=name + 'dwconv_pad')(x)
conv_pad = 'valid'
else:
conv_pad = 'same'
x = tf.keras.layers.DepthwiseConv2D(
kernel_size,
strides=strides,
padding=conv_pad,
use_bias=False,
depthwise_initializer=CONV_KERNEL_INITIALIZER,
name=name + 'dwconv')(x)
x = build_normalization(**self.normalization, name=name + 'bn')(x)
x = build_activation(**self.activation, name=name + 'activation')(x)
# Squeeze and Excitation phase
if 0 < se_ratio <= 1:
filters_se = max(1, int(filters_in * se_ratio))
se = tf.keras.layers.GlobalAveragePooling2D(name=name +
'se_squeeze')(x)
se = tf.keras.layers.Reshape((1, 1, filters),
name=name + 'se_reshape')(se)
se = tf.keras.layers.Conv2D(
filters_se,
1,
padding='same',
activation=activation["activation"],
kernel_initializer=CONV_KERNEL_INITIALIZER,
name=name + 'se_reduce')(se)
se = tf.keras.layers.Conv2D(
filters,
1,
padding='same',
activation='sigmoid',
kernel_initializer=CONV_KERNEL_INITIALIZER,
name=name + 'se_expand')(se)
x = tf.keras.layers.multiply([x, se], name=name + 'se_excite')
# Output phase
x = tf.keras.layers.Conv2D(filters_out,
1,
padding='same',
use_bias=False,
kernel_initializer=CONV_KERNEL_INITIALIZER,
name=name + 'project_conv')(x)
x = build_normalization(**self.normalization,
name=name + 'project_bn')(x)
if id_skip and strides == 1 and filters_in == filters_out:
if drop_rate > 0:
x = tf.keras.layers.Dropout(drop_rate,
noise_shape=(None, 1, 1, 1),
name=name + 'drop')(x)
x = tf.keras.layers.add([x, inputs], name=name + 'add')
return x
def build_model(self):
inputs = tf.keras.layers.Input((None, None, 3))
trainable = 1 not in self.frozen_stages
x = tf.keras.layers.Lambda(function=lambda inp: (
(inp * (1. / 255.)) - self._rgb_mean) / self._rgb_std)(inputs)
x = tf.keras.layers.ZeroPadding2D(padding=((3, 3), (3, 3)))(x)
x = build_convolution(self.convolution,
filters=64,
kernel_size=7,
strides=self.strides[0],
use_bias=False,
trainable=trainable,
name="conv1/conv")(x)
x = build_normalization(**self.normalization, name="conv1/bn")(x)
x1 = tf.keras.layers.Activation(**self.activation,
name="conv1/relu")(x)
trainable = 2 not in self.frozen_stages
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(x)
x = tf.keras.layers.MaxPool2D(3, strides=self.strides[1],
name="pool1")(x)
x = self.dense_block(x,
blocks=self.blocks[0],
dilation_rate=self.dilation_rates[1],
drop_rate=self.drop_rate,
trainable=trainable,
name="conv2")
trainable = 3 not in self.frozen_stages
preact2, x = self.transition_block(x, 0.5, self.strides[2], trainable,
"pool2")
x = self.dense_block(x,
blocks=self.blocks[1],
dilation_rate=self.dilation_rates[2],
drop_rate=self.drop_rate,
trainable=trainable,
name="conv3")
trainable = 4 not in self.frozen_stages
preact3, x = self.transition_block(x, 0.5, self.strides[3], trainable,
"pool3")
x = self.dense_block(x,
blocks=self.blocks[2],
dilation_rate=self.dilation_rates[3],
drop_rate=self.drop_rate,
trainable=trainable,
name="conv4")
trainable = 5 not in self.frozen_stages
preact4, x = self.transition_block(x, 0.5, self.strides[4], trainable,
"pool4")
x = self.dense_block(x,
blocks=self.blocks[3],
dilation_rate=self.dilation_rates[4],
drop_rate=self.drop_rate,
trainable=trainable,
name="conv5")
x = build_normalization(**self.normalization, name="bn")(x)
x5 = tf.keras.layers.Activation(**self.activation, name="relu")(x)
if self._is_classifier:
x = tf.keras.layers.GlobalAvgPool2D(name="avg_pool")(x5)
x = tf.keras.layers.Dense(self.num_classes, name="fc1000")(x)
return tf.keras.Model(inputs=inputs, outputs=x)
else:
outputs = [
o for i, o in enumerate([x1, preact2, preact3, preact4, x5])
if i + 1 in self.output_indices
]
return tf.keras.Model(inputs=inputs,
outputs=outputs,
name=self.name)
def bottleneck_v1(x,
convolution,
filters,
strides=1,
dilation_rate=1,
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
weight_decay=0.,
dropblock=None,
use_conv_shortcut=True,
name=None):
"""A residual block.
Args:
x: input tensor.
filters: integer, filters of the bottleneck layer.
convolution: The convolution type.
strides: default 1, stride of the first layer.
dilation_rate: default 1, dilation rate in 3x3 convolution.
activation: the activation layer name.
trainable: does this block is trainable.
normalization: the normalization, e.g. "batch_norm", "group_norm" etc.
weight_decay: weight decay.
dropblock: the arguments in DropBlock2D
use_conv_shortcut: default True, use convolution shortcut if True,
otherwise identity shortcut.
name: string, block label.
Returns:
Output tensor for the residual block.
"""
if use_conv_shortcut is True:
shortcut = build_convolution(
convolution,
filters=4 * filters,
kernel_size=1,
strides=strides,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_0_conv")(x)
shortcut = build_normalization(**normalization,
name=name + "_0_bn")(shortcut)
else:
shortcut = x
if dropblock is not None:
shortcut = DropBlock2D(**dropblock,
name=name + "_0_dropblock")(shortcut)
x = build_convolution(
convolution,
filters=filters,
kernel_size=1,
strides=strides,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_1_conv")(x)
x = build_normalization(**normalization, name=name + "_1_bn")(x)
x = tf.keras.layers.Activation(**activation, name=name + "_1_relu")(x)
if dropblock is not None:
x = DropBlock2D(**dropblock, name=name + "_1_dropblock")(x)
x = build_convolution(
convolution,
filters=filters,
kernel_size=3,
strides=1,
padding="SAME",
dilation_rate=dilation_rate,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_2_conv")(x)
x = build_normalization(**normalization, name=name + "_2_bn")(x)
x = tf.keras.layers.Activation(**activation, name=name + "_2_relu")(x)
if dropblock is not None:
x = DropBlock2D(**dropblock, name=name + "_2_dropblock")(x)
x = build_convolution(
convolution,
filters=4 * filters,
kernel_size=1,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_3_conv")(x)
x = build_normalization(**normalization, name=name + "_3_bn")(x)
if dropblock is not None:
x = DropBlock2D(**dropblock, name=name + "_3_dropblock")(x)
x = tf.keras.layers.Add(name=name + "_add")([shortcut, x])
x = tf.keras.layers.Activation(**activation, name=name + "_out")(x)
return x
def bottleneckx(x,
convolution,
filters,
strides=1,
dilation_rate=1,
cardinality=32,
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
weight_decay=0.,
use_conv_shortcut=True,
name=None):
"""A residual block.
Args:
x: input tensor.
filters: integer, filters of the bottleneck layer.
convolution: The convolution type.
strides: default 1, stride of the first layer.
dilation_rate: default 1, dilation rate in 3x3 convolution.
cardinality: default 32, the cardinality in resnext.
activation: the activation layer name.
trainable: does this block is trainable.
normalization: the normalization, e.g. "batch_norm", "group_norm" etc.
weight_decay: the weight decay.
use_conv_shortcut: default True, use convolution shortcut if True,
otherwise identity shortcut.
name: string, block label.
Returns:
Output tensor for the residual block.
"""
bn_axis = 3 if tf.keras.backend.image_data_format(
) == "channels_last" else 1
if use_conv_shortcut:
shortcut = build_convolution(
convolution,
filters=(64 // cardinality) * filters,
kernel_size=1,
strides=strides,
use_bias=False,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_0_conv")(x)
shortcut = build_normalization(**normalization,
name=name + "_0_bn")(shortcut)
else:
shortcut = x
x = tf.keras.layers.Conv2D(filters,
1,
use_bias=False,
trainable=trainable,
name=name + "_1_conv")(x)
x = build_normalization(**normalization, name=name + "_1_bn")(x)
x = tf.keras.layers.Activation(activation, name=name + "_1_relu")(x)
c = filters // cardinality
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)),
name=name + "_2_pad")(x)
x = tf.keras.layers.DepthwiseConv2D(3,
strides=strides,
depth_multiplier=c,
dilation_rate=dilation_rate,
use_bias=False,
trainable=trainable,
name=name + "_2_conv")(x)
kernel = np.zeros((1, 1, filters * c, filters), dtype=np.float32)
for i in range(filters):
start = (i // c) * c * c + i % c
end = start + c * c
kernel[:, :, start:end:c, i] = 1.
x = tf.keras.layers.Conv2D(filters,
1,
use_bias=False,
trainable=False,
kernel_initializer={
"class_name": "Constant",
"config": {
"value": kernel
}
},
name=name + "_2_gconv")(x)
x = build_normalization(**normalization, name=name + "_2_bn")(x)
x = tf.keras.layers.Activation("relu", name=name + "_2_relu")(x)
x = tf.keras.layers.Conv2D((64 // cardinality) * filters,
1,
use_bias=False,
trainable=trainable,
name=name + "_3_conv")(x)
x = build_normalization(**normalization, name=name + "_3_bn")(x)
x = tf.keras.layers.Add(name=name + "_add")([shortcut, x])
x = tf.keras.layers.Activation("relu", name=name + "_out")(x)
return x
def bottleneck_v2(x,
convolution,
filters,
strides=1,
dilation_rate=1,
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
weight_decay=0.,
dropblock=None,
use_conv_shortcut=True,
name=None):
"""A residual block.
Args:
x: input tensor.
filters: integer, filters of the bottleneck layer.
convolution: The convolution type.
strides: default 1, stride of the first layer.
dilation_rate: default 1, dilation rate in 3x3 convolution.
activation: the activation layer name.
trainable: does this block is trainable.
normalization: the normalization, e.g. "batch_norm", "group_norm" etc.
weight_decay: weight decay.
dropblock: the arguments in DropBlock2D.
use_conv_shortcut: default True, use convolution shortcut if True,
otherwise identity shortcut.
name: string, block label.
Returns:
Output tensor for the residual block.
"""
bn_axis = 3 if tf.keras.backend.image_data_format(
) == "channels_last" else 1
preact = build_normalization(**normalization, name=name + "_preact_bn")(x)
preact = tf.keras.layers.Activation(**activation,
name=name + "_preact_relu")(preact)
if use_conv_shortcut is True:
shortcut = build_convolution(
convolution,
filters=4 * filters,
kernel_size=1,
strides=strides,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_0_conv")(preact)
else:
shortcut = tf.keras.layers.MaxPooling2D(
1, strides=strides)(x) if strides > 1 else x
if dropblock is not None:
shortcut = DropBlock2D(**dropblock,
name=name + "_0_dropblock")(shortcut)
x = build_convolution(
convolution,
filters=filters,
kernel_size=1,
strides=1,
use_bias=False,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_1_conv")(preact)
x = build_normalization(**normalization, name=name + "_1_bn")(x)
x = tf.keras.layers.Activation(**activation, name=name + "_1_relu")(x)
if dropblock is not None:
x = DropBlock2D(**dropblock, name=name + "_1_dropblock")(x)
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)),
name=name + "_2_pad")(x)
x = build_convolution(
convolution,
filters=filters,
kernel_size=3,
strides=strides,
dilation_rate=1 if strides > 1 else dilation_rate,
use_bias=False,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_2_conv")(x)
x = build_normalization(**normalization, name=name + "_2_bn")(x)
x = tf.keras.layers.Activation(**activation, name=name + "_2_relu")(x)
if dropblock is not None:
x = DropBlock2D(**dropblock)(x)
x = build_convolution(
convolution,
filters=4 * filters,
kernel_size=1,
trainable=trainable,
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
name=name + "_3_conv")(x)
if dropblock is not None:
x = DropBlock2D(**dropblock, name=name + "_2_dropblock")(x)
x = tf.keras.layers.Add(name=name + "_out")([shortcut, x])
return x, preact
def __init__(self,
filters,
kernel_size,
strides=1,
groups=1,
padding="same",
data_format="channels_last",
dilation_rate=1,
kernel_initializer="he_normal",
trainable=True,
use_bias=False,
normalization=dict(normalization="batch_norm",
axis=-1,
trainable=True),
gamma_zeros=False,
activation=None,
dropblock=None,
name=None):
super(ConvNormActBlock, self).__init__(name=name)
self.strides = strides
if isinstance(kernel_size, int):
kernel_size = (kernel_size, kernel_size)
if isinstance(strides, int):
strides = (strides, strides)
# if strides != (1, 1):
# dilation_rate = (1, 1)
if isinstance(dilation_rate, int):
dilation_rate = (dilation_rate, dilation_rate)
if strides == (1, 1):
p = 1
padding == "same"
else:
p = ((kernel_size[0] - 1) // 2 * dilation_rate[0],
(kernel_size[1] - 1) * dilation_rate[1] // 2)
self.pad = tf.keras.layers.ZeroPadding2D(p,
data_format=data_format)
padding = "valid"
self.conv = tf.keras.layers.Conv2D(
filters=filters,
kernel_size=kernel_size,
strides=strides,
padding=padding,
data_format=data_format,
dilation_rate=dilation_rate,
groups=groups,
use_bias=normalization is None or use_bias,
trainable=trainable,
kernel_initializer=kernel_initializer,
name="conv2d")
channel_axis = -1 if data_format == "channels_last" else 1
if not trainable and normalization is not None:
normalization["trainable"] = False
if normalization is not None:
normalization["axis"] = channel_axis
norm_name = ("batch_norm"
if "batch_norm" in normalization["normalization"] else
normalization["normalization"])
if norm_name == "batch_norm" and not trainable:
normalization["normalization"] = "frozen_batch_norm"
self.norm = build_normalization(
gamma_initializer="zeros" if gamma_zeros else "ones",
name=norm_name,
**normalization)
else:
self.norm = None
# self.act = build_activation(**activation, name=activation["activation"]) if activation is not None else None
# self.dropblock = DropBlock2D(**dropblock, name="dropblock") if dropblock is not None else None
if activation is not None:
self.act = (build_activation(**activation) if isinstance(
activation, dict) else build_activation(activation=activation))
else:
self.act = None
def build_bifpn_layer(feats,
input_size,
feat_dims,
convolution="separable_conv2d",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="swish"),
min_level=3,
max_level=7,
pool_type=None,
apply_bn=True,
fusion_type="weighted_sum",
name="fpn_cells",
**kwargs):
F = lambda x: 1.0 / (2**x) # Resolution size for a given feature level.
node_configs = [
{
"width_ratio": F(6),
"inputs_offsets": [3, 4]
},
{
"width_ratio": F(5),
"inputs_offsets": [2, 5]
},
{
"width_ratio": F(4),
"inputs_offsets": [1, 6]
},
{
"width_ratio": F(3),
"inputs_offsets": [0, 7]
},
{
"width_ratio": F(4),
"inputs_offsets": [1, 7, 8]
},
{
"width_ratio": F(5),
"inputs_offsets": [2, 6, 9]
},
{
"width_ratio": F(6),
"inputs_offsets": [3, 5, 10]
},
{
"width_ratio": F(7),
"inputs_offsets": [4, 11]
},
]
num_output_connections = [0 for _ in feats]
for i, fnode in enumerate(node_configs):
nodes = []
node_name = name + "/fnode{}".format(i)
new_node_width = int(fnode["width_ratio"] * input_size[1])
for idx, input_offset in enumerate(fnode["inputs_offsets"]):
input_node = feats[input_offset]
num_output_connections[input_offset] += 1
input_node = resample_feature_map(
input_node,
new_node_width,
feat_dims,
convolution=convolution,
normalization=normalization,
activation=activation,
pool_type=pool_type,
apply_bn=apply_bn,
name=node_name +
"/resample_{}_{}_{}".format(idx, input_offset, len(feats)))
nodes.append(input_node)
if fusion_type == "weighted_sum":
if len(fnode["inputs_offsets"]) == 2:
new_node = WeightedFusion2(name=node_name)(nodes)
if len(fnode["inputs_offsets"]) == 3:
new_node = WeightedFusion3(name=node_name)(nodes)
elif fusion_type == "sum":
new_node = tf.keras.layers.Add(name=node_name)(nodes)
else:
raise ValueError("Unknown fusion type: {}".format(fusion_type))
new_node_name = node_name + "/op_after_combine{}".format(len(feats))
new_node = build_activation(**activation,
name=new_node_name + "/" +
activation["activation"])(new_node)
new_node = build_convolution(convolution,
filters=feat_dims,
kernel_size=(3, 3),
padding="same",
name=new_node_name + "/conv")(new_node)
new_node = build_normalization(**normalization,
name=new_node_name + "/bn")(new_node)
# new_node = build_activation(
# **activation, name=new_node_name + "/" + activation["activation"] + "_1")(new_node)
feats.append(new_node)
num_output_connections.append(0)
outputs = []
for level in range(min_level, max_level + 1):
for i, fnode in enumerate(reversed(node_configs)):
if fnode["width_ratio"] == F(level):
outputs.append(feats[-1 - i])
break
return outputs
def __init__(self,
convolution="conv2d",
normalization="batch_norm",
activation="relu",
output_indices=(-1, ),
strides=(2, 2, 2, 2, 2),
dilation_rates=(1, 1, 1, 1, 1),
freezing_stages=(-1, ),
freezing_batch_normalization=False,
group=32,
weight_decay=0.):
super(NewResNet18, self).__init__(
convolution=convolution,
normalization=normalization,
activation=activation,
output_indices=output_indices,
strides=strides,
dilation_rates=dilation_rates,
freezing_stages=freezing_stages,
freezing_batch_normalization=freezing_batch_normalization,
weight_decay=weight_decay)
axis = 3 if tf.keras.backend.image_data_format(
) == "channels_last" else 1
self.conv1_1 = tf.keras.Sequential([
build_convolution(
convolution,
filters=16,
kernsl_size=7,
padding="same",
dilation_rate=dilation_rates[0],
kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
use_bias=False),
build_normalization(normalization, asix=axis, group=group),
tf.keras.layers.Activation(activation)
])
self.conv1_2 = BasicBlock(convolution=convolution,
filters=16,
strides=1,
dilation_rate=dilation_rates[0],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=False)
self.conv1_3 = BasicBlock(convolution=convolution,
filters=32,
strides=strides[0],
dilation_rate=dilation_rates[0],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=True)
self.conv2_1 = BasicBlock(convolution=convolution,
filters=64,
strides=strides[1],
dilation_rate=dilation_rates[1],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=True)
self.conv2_2 = BasicBlock(convolution=convolution,
filters=64,
strides=1,
dilation_rate=dilation_rates[1],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=False)
self.conv3_1 = BasicBlock(convolution=convolution,
filters=128,
strides=strides[2],
dilation_rate=dilation_rates[2],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=True)
self.conv3_2 = BasicBlock(convolution=convolution,
filters=128,
strides=1,
dilation_rate=dilation_rates[2],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=False)
self.conv4_1 = BasicBlock(convolution=convolution,
filters=256,
strides=strides[3],
dilation_rate=dilation_rates[3],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=True)
self.conv4_2 = BasicBlock(convolution=convolution,
filters=256,
strides=1,
dilation_rate=dilation_rates[3],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=False)
self.conv5_1 = BasicBlock(convolution=convolution,
filters=512,
strides=strides[4],
dilation_rate=dilation_rates[4],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=True)
self.conv5_2 = BasicBlock(convolution=convolution,
filters=512,
strides=1,
dilation_rate=dilation_rates[4],
normalization=normalization,
group=16,
activation=activation,
weight_decay=weight_decay,
use_conv_shortcut=False)
def resample_feature_map(feat,
target_width,
target_num_channels,
convolution="separable_conv2d",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="swish"),
pool_type=None,
apply_bn=True,
name="resample"):
"""Resample input feature map to have target number of channels and width."""
_, _, width, num_channels = tf.keras.backend.int_shape(feat)
if width > target_width:
if num_channels != target_num_channels:
feat = build_convolution("conv2d",
filters=target_num_channels,
kernel_size=(1, 1),
padding="same",
name=name + "/conv2d")(feat)
if apply_bn:
feat = build_normalization(**normalization,
name=name + "/bn")(feat)
strides = int(width // target_width)
# if strides >= 2:
# feat = tf.keras.layers.ZeroPadding2D(
# padding=imagenet_utils.correct_pad(feat, strides + 1),
# name=name + '/stem/conv_pad')(feat)
# pad = "valid"
# else:
# pad = "same"
if pool_type == "max" or pool_type is None:
feat = tf.keras.layers.MaxPool2D(pool_size=strides + 1,
strides=[strides, strides],
padding="same",
name=name + "/max_pool")(feat)
elif pool_type == "avg":
feat = tf.keras.layers.AvgPool2D(pool_size=strides + 1,
strides=[strides, strides],
padding="same",
name=name + "/avg_pool")(feat)
else:
raise ValueError("Unknown pooling type: {}".format(pool_type))
else:
if num_channels != target_num_channels:
feat = build_convolution("conv2d",
filters=target_num_channels,
kernel_size=(1, 1),
padding="same",
name=name + "/conv2d")(feat)
if apply_bn:
feat = build_normalization(**normalization,
name=name + "/bn")(feat)
if width < target_width:
scale = target_width // width
feat = NearestUpsampling2D(scale=scale,
name=name + "/nearest_upsampling")(feat)
return feat
