def bottleneck(inputs,
shortcut,
filters,
strides=1,
dilation_rate=1,
expansion=2,
cardinality=32,
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
kernel_initializer="he_normal",
name="bottleneck"):
x = ConvNormActBlock(filters=filters // expansion,
kernel_size=1,
strides=1,
data_format=data_format,
dilation_rate=1,
trainable=trainable,
normalization=normalization,
activation=activation,
kernel_initializer=kernel_initializer,
name=name + "/conv1")(inputs)
x = ConvNormActBlock(filters=filters // expansion,
kernel_size=3,
strides=strides,
data_format=data_format,
dilation_rate=dilation_rate if strides == 1 else 1,
trainable=trainable,
normalization=normalization,
kernel_initializer=kernel_initializer,
activation=activation,
name=name + "/conv2")(x)
x = ConvNormActBlock(filters=filters,
kernel_size=1,
strides=1,
data_format=data_format,
dilation_rate=1,
trainable=trainable,
normalization=normalization,
activation=None,
kernel_initializer=kernel_initializer,
name=name + "/conv3")(x)
if shortcut is None:
shortcut = inputs
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
if isinstance(activation, dict):
act_name = activation["activation"]
x = build_activation(name=name + "/" + act_name, **activation)(x)
else:
x = build_activation(activation=activation,
name=name + "/" + activation)(x)
return x
def root(inputs,
filters,
kernel_size,
residual,
trainable=True,
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",
name="root"):
channel_axis = -1 if data_format == "channels_last" else 1
x = tf.keras.layers.Concatenate(axis=channel_axis, name=name + "/cat")(inputs)
x = ConvNormActBlock(filters=filters,
kernel_size=kernel_size,
strides=1,
data_format=data_format,
normalization=normalization,
activation=None,
trainable=trainable,
kernel_initializer=kernel_initializer,
name=name + "/conv")(x)
if residual:
x = tf.keras.layers.Add(name=name + "/add")([x, inputs[0]])
if isinstance(activation, dict):
act_name = activation["activation"]
x = build_activation(name=name + "/" + act_name, **activation)(x)
else:
x = build_activation(activation=activation, name=name + "/" + activation)(x)
return x
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 bottleneck(inputs,
filters,
strides=1,
dilation_rate=1,
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="leaky_relu", alpha=0.1),
kernel_initializer="glorot_uniform",
trainable=True,
shortcut=True,
name="bottleneck"):
channel_axis = -1 if data_format == "channels_last" else 1
normalization["axis"] = channel_axis
if not trainable:
normalization["trainable"] = trainable
x = ConvNormActBlock(filters=filters,
kernel_size=3,
strides=strides,
dilation_rate=1 if strides == 2 else dilation_rate,
normalization=normalization,
activation=activation,
data_format=data_format,
trainable=trainable,
kernel_initializer=kernel_initializer,
name=name + "/conv1")(inputs)
x = ConvNormActBlock(filters=filters,
kernel_size=3,
dilation_rate=dilation_rate,
data_format=data_format,
normalization=normalization,
activation=None,
trainable=trainable,
kernel_initializer=kernel_initializer,
name=name + "/conv2")(x)
shortcut = inputs
if tf.keras.backend.int_shape(shortcut)[-1] != filters or strides != 1:
shortcut = ConvNormActBlock(filters=filters,
kernel_size=1,
strides=strides,
data_format=data_format,
kernel_initializer=kernel_initializer,
trainable=trainable,
activation=None,
normalization=normalization,
name=name + "/skip")(shortcut)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
if isinstance(activation, str):
x = build_activation(activation=activation,
name=name + "/" + activation)(x)
else:
x = build_activation(name=name + "/" + activation["activation"],
**activation)(x)
return x
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 __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 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 _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 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(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 __init__(self, cfg, **kwargs):
super(OneNetHead, self).__init__(cfg, **kwargs)
assert cfg.assigner.assigner == "MinCostAssigner"
self.feat = tf.keras.layers.Conv2D(
filters=cfg.feat_dims,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="feat1")
self.classifier = tf.keras.layers.Conv2D(
filters=self.num_classes,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
bias_initializer=tf.keras.initializers.Constant(
-math.log((1. - cfg.prior) / cfg.prior)),
name="predicted_class")
self.regressor = tf.keras.layers.Conv2D(
filters=4,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="predicted_box")
self.act = build_activation(**cfg.activation.as_dict())
self.assigner = build_assigner(**cfg.assigner.as_dict()) if cfg.get("assigner") is not None else None
self.bbox_loss_func = build_loss(**cfg.bbox_loss.as_dict()) if cfg.get("bbox_loss") is not None else None
self._use_iou_loss = False
if self.bbox_loss_func is not None:
self._use_iou_loss = "IoU" in cfg.bbox_loss.loss
self.label_loss_func = build_loss(**cfg.label_loss.as_dict()) if cfg.get("label_loss") is not None else None
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 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 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 bottleneck_v1b(inputs,
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,
dropblock=None,
last_gamma_init_zero=False,
use_conv_shortcut=False,
data_format="channels_last",
expansion=4,
name="bottleneck_v1b"):
"""A residual block.
Args:
filters: integer, filters of the bottleneck layer.
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.
dropblock: the arguments in DropBlock2D
use_conv_shortcut: default True, use convolution shortcut if True,
otherwise identity shortcut.
Returns:
Output tensor for the residual block.
"""
x = ConvNormActBlock(filters=filters,
kernel_size=1,
trainable=trainable,
normalization=normalization,
activation=activation,
data_format=data_format,
dropblock=dropblock,
name=name + "/conv1")(inputs)
x = ConvNormActBlock(filters=filters,
kernel_size=3,
strides=strides,
dilation_rate=1 if strides == 2 else dilation_rate,
data_format=data_format,
trainable=trainable,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/conv2")(x)
x = ConvNormActBlock(filters=filters * expansion,
kernel_size=1,
data_format=data_format,
trainable=trainable,
normalization=normalization,
gamma_zeros=last_gamma_init_zero,
activation=None,
name=name + "/conv3")(x)
shortcut = inputs
if use_conv_shortcut:
if dilation_rate == 1:
shortcut = tf.keras.layers.AvgPool2D(pool_size=strides,
strides=strides,
padding="same",
data_format=data_format,
name=name + "/avgpool")(shortcut)
else:
shortcut = tf.keras.layers.AvgPool2D(pool_size=1,
strides=1,
padding="same",
data_format=data_format,
name=name + "/avgpool")(shortcut)
shortcut = ConvNormActBlock(filters=filters * expansion,
kernel_size=1,
data_format=data_format,
trainable=trainable,
normalization=normalization,
dropblock=dropblock,
activation=None,
name=name + "/downsample")(shortcut)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
x = build_activation(name=name + "/" + activation["activation"], **activation)(x)
return x
def bottle2neck(inputs,
filters,
strides=1,
scale=4,
base_width=26,
dilation_rate=1,
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
downsample=False,
trainable=True,
dropblock=None,
stype="normal",
expansion=4,
name="Bottle2neck"):
width = int(math.floor(filters * (base_width / 64.)))
channel_axis = -1 if data_format == "channels_last" else 1
x = ConvNormActBlock(filters=width * scale,
kernel_size=1,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/conv1")(inputs)
num_convs = scale if scale == 1 else scale - 1
spx = tf.keras.layers.Lambda(
lambda inp: tf.split(inp, scale, channel_axis),
name=name + "/split")(x)
for i in range(num_convs):
if i == 0 or stype == "stage":
sp = spx[i]
else:
sp = tf.keras.layers.Add(name=name + "/add%d" % i)([sp, spx[i]])
sp = ConvNormActBlock(
filters=width,
kernel_size=3,
strides=strides,
data_format=data_format,
dilation_rate=dilation_rate if strides == 1 else 1,
trainable=trainable,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/convs/%d" % i)(sp)
if i == 0:
x = sp
else:
x = tf.keras.layers.Concatenate(channel_axis,
name=name + "/cat%d" % i)([x, sp])
if scale != 1 and stype == "normal":
x = tf.keras.layers.Concatenate(channel_axis,
name=name + "/cat%d" % num_convs)(
[x, spx[num_convs]])
elif scale != 1 and stype == "stage":
padding = "same"
sp_ = spx[num_convs]
if strides != 1:
sp_ = tf.keras.layers.ZeroPadding2D(((1, 1), (1, 1)),
data_format=self.data_format,
name=name + "/pad")(sp_)
padding = "valid"
sp_ = tf.keras.layers.AvgPool2D(3,
strides,
padding,
data_format,
name=name + "/avgpool")(sp_)
x = tf.keras.layers.Concatenate(channel_axis,
name=name +
"/cat%d" % num_convs)([x, sp_])
x = ConvNormActBlock(filters=filters * expansion,
kernel_size=1,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/conv3")(x)
shortcut = inputs
if downsample:
shortcut = ConvNormActBlock(filters=filters * expansion,
kernel_size=1,
strides=strides,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/downsample")(shortcut)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
x = build_activation(**activation,
name=name + "/" + activation["activation"])(x)
return x
def bottleneckx(inputs,
filters,
cardinality,
bottleneck_width,
strides=1,
dilation_rate=1,
use_se=False,
normalization=dict(normalization="batch_norm", momentum=0.9, epsilon=1e-5, axis=-1, trainable=True),
activation=dict(activation="relu"),
trainable=True,
downsample=False,
dropblock=None,
avg_down=False,
last_gamma=False,
data_format="channels_last",
expansion=4,
name=None):
d = int(math.floor(filters * (bottleneck_width / 64)))
group_width = cardinality * d
x = ConvNormActBlock(filters=group_width,
kernel_size=1,
strides=1,
trainable=trainable,
data_format=data_format,
dropblock=dropblock,
normalization=normalization,
activation=activation,
name=name + "/conv1")(inputs)
x = ConvNormActBlock(filters=group_width,
kernel_size=(3, 3),
strides=strides,
groups=cardinality,
dilation_rate=1 if strides == 1 else dilation_rate,
data_format=data_format,
trainable=trainable,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/conv2")(x)
x = ConvNormActBlock(filters=filters * expansion,
kernel_size=1,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/conv3")(x)
if use_se:
pool_axis = [1, 2] if data_format == "channels_last" else [2, 3]
se = tf.keras.layers.Lambda(lambda inp: tf.reduce_mean(inp, pool_axis, keepdims=True),
name=name + "/se/pool")(x)
se = tf.keras.layers.Conv2D(filters=filters // 4,
kernel_size=(1, 1),
data_format=data_format,
name=name + "/conv4")(se)
se = tf.keras.layers.Conv2D(filters=filters * expansion,
kernel_size=(1, 1),
data_format=data_format,
name=name + "/conv5")(se)
se = tf.keras.layers.Activation("sigmoid", name=name + "/sigmoid")(se)
x = tf.keras.layers.Multiply(name=name + "/multiply")([x, se])
shortcut = inputs
if downsample:
if avg_down:
shortcut = tf.keras.layers.AvgPool2D(pool_size=strides,
strides=strides,
padding="same",
data_format=data_format,
name=name + "/shorcut/avg_pool")(shortcut)
shortcut = ConvNormActBlock(kernel_size=1,
filters=filters * expansion,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/shortcut")(shortcut)
else:
shortcut = ConvNormActBlock(kernel_size=1,
strides=strides,
filters=filters * expansion,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/shortcut")(shortcut)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
x = build_activation(**activation, name=name + "/" + activation["activation"])(x)
return x
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,
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 shuffle_v1_block(inputs,
in_filters,
out_filters,
kernel_size,
mid_filters,
strides=1,
groups=1,
first_group=False,
kernel_initializer="he_normal",
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
name="shuffle_v1_block"):
x = ConvNormActBlock(filters=mid_filters,
kernel_size=1,
strides=1,
groups=1 if first_group else groups,
kernel_initializer=kernel_initializer,
normalization=normalization,
activation=activation,
trainable=trainable,
name=name + "/conv1")(inputs)
x = DepthwiseConvNormActBlock(kernel_size=kernel_size,
strides=strides,
kernel_initializer=kernel_initializer,
normalization=normalization,
activation=None,
trainable=trainable,
name=name + "/conv2")(x)
if groups > 1:
x = ChannelShuffle(groups=groups,
data_format=data_format,
name=name + "/channel_shuffle")(x)
x = ConvNormActBlock(filters=out_filters -
in_filters if strides == 2 else out_filters,
kernel_size=1,
strides=1,
groups=groups,
kernel_initializer=kernel_initializer,
normalization=normalization,
activation=None,
trainable=trainable,
name=name + "/conv3")(x)
if strides == 1:
x = tf.keras.layers.Add(name=name + "/sum")([x, inputs])
return x
inputs = tf.keras.layers.AvgPool2D(pool_size=3,
strides=strides,
padding="same",
data_format=data_format,
name=name + "/avgpool")(inputs)
x = build_activation(name=name + "/%s" % activation["activation"],
**activation)(x)
axis = -1 if data_format == "channels_last" else 1
return tf.keras.layers.Concatenate(axis=axis,
name=name + "/concat")([inputs, x])
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 bottleneck(inputs,
filters,
strides=1,
dilation_rate=1,
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
dropblock=None,
use_conv_shortcut=True,
strides_in_1x1=False,
expansion=4,
name=None):
"""A residual block.
Args:
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.
data_format: default channels_last,
activation: the activation layer name.
trainable: does this block is trainable.
normalization: the normalization, e.g. "batch_norm", "group_norm" etc.
dropblock: the arguments in DropBlock2D
use_conv_shortcut: default True, use convolution shortcut if True,
otherwise identity shortcut.
name: string, block label.
"""
strides_1x1, strides_3x3 = (strides, 1) if strides_in_1x1 else (1, strides)
x = ConvNormActBlock(filters=filters,
kernel_size=1,
strides=strides_1x1,
trainable=trainable,
dropblock=dropblock,
data_format=data_format,
normalization=normalization,
activation=activation,
name=name + "/conv1")(inputs)
x = ConvNormActBlock(
filters=filters,
kernel_size=3,
strides=strides_3x3,
dilation_rate=dilation_rate if strides == 1 or strides_1x1 else 1,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=activation,
name=name + "/conv2")(x)
x = ConvNormActBlock(filters=expansion * filters,
kernel_size=1,
trainable=trainable,
data_format=data_format,
normalization=normalization,
activation=None,
name=name + "/conv3")(x)
shortcut = inputs
if use_conv_shortcut is True:
shortcut = ConvNormActBlock(filters=expansion * filters,
kernel_size=1,
strides=strides,
data_format=data_format,
trainable=trainable,
dropblock=dropblock,
normalization=normalization,
activation=None,
name=name + "/shortcut")(shortcut)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
x = build_activation(**activation,
name=name + "/" + activation["activation"])(x)
return x
def linear_bottleneck_v1(inputs,
t,
in_filters,
out_filters,
strides=1,
use_se=True,
se_ratio=12,
kernel_initializer="he_normal",
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
trainable=True,
name="linear_bottleneck"):
use_shorcut = strides == 1 and in_filters <= out_filters
x = inputs
if t != 1:
dw_filters = in_filters * t
x = ConvNormActBlock(filters=dw_filters,
kernel_size=1,
normalization=normalization,
activation=dict(activation="swish"),
data_format=data_format,
kernel_initializer=kernel_initializer,
trainable=trainable,
name=name + "/conv1")(x)
else:
dw_filters = in_filters
x = DepthwiseConvNormActBlock(kernel_size=3,
strides=strides,
data_format=data_format,
kernel_initializer=kernel_initializer,
normalization=normalization,
name=name + "/conv2")(x)
if use_se:
x = squeeze_excitation(inputs=x,
in_filters=dw_filters,
se_ratio=se_ratio,
data_format=data_format,
trainable=trainable,
normalization=normalization,
name=name + "/se")
x = build_activation(activation="relu6", name=name + "/relu6")(x)
x = ConvNormActBlock(filters=out_filters,
kernel_size=1,
data_format=data_format,
kernel_initializer=kernel_initializer,
trainable=trainable,
normalization=normalization,
activation=None,
name=name + "/conv3")(x)
if use_shorcut:
x = Shortcut(in_filters=in_filters,
data_format=data_format,
name=name + "/shortcut")(x, inputs)
return x
def build_model(self):
trainable = 0 not in self.frozen_stages
def _fn(inp):
# if self.data_format == "channels_first":
# inp = inp[:, ::-1, :, :]
# else:
# inp = inp[:, :, :, ::-1]
mean = tf.constant([[[[0.408, 0.447, 0.470]]]], inp.dtype) * 255.
std = 1. / (tf.constant([[[[0.289, 0.274, 0.278]]]], inp.dtype) *
255.)
return (inp - mean) * std
curr_filters = self.filters[0]
kernel_initializer = "he_normal"
inputs = tf.keras.layers.Lambda(function=_fn,
name="norm_input")(self.img_input)
x = ConvNormActBlock(filters=128,
kernel_size=7,
strides=2,
normalization=self.normalization,
activation=self.activation,
trainable=trainable,
data_format=self.data_format,
kernel_initializer=kernel_initializer,
name="pre/0")(inputs)
inter = bottleneck(inputs=x,
filters=256,
strides=2,
data_format=self.data_format,
kernel_initializer=kernel_initializer,
normalization=self.normalization,
activation=self.activation,
trainable=1 not in self.frozen_stages,
name="pre/1")
self.prev_filters = 256
outputs = []
for i in range(self.num_stacks):
kp = self.stack_module(inter,
self._n,
self.filters,
self.blocks,
name="kps/%d" % i)
cnv = ConvNormActBlock(filters=curr_filters,
kernel_size=3,
data_format=self.data_format,
normalization=self.normalization,
activation=self.activation,
name="cnvs/%d" % i)(kp)
outputs.append(cnv)
if i < self.num_stacks - 1:
x1 = ConvNormActBlock(filters=curr_filters,
kernel_size=1,
data_format=self.data_format,
normalization=self.normalization,
activation=None,
name="inters_/%d" % i)(inter)
x2 = ConvNormActBlock(filters=curr_filters,
kernel_size=1,
data_format=self.data_format,
normalization=self.normalization,
activation=None,
name="cnvs_/%d" % i)(cnv)
inter = tf.keras.layers.Add()([x1, x2])
inter = build_activation(**self.activation)(inter)
inter = bottleneck(inter,
filters=curr_filters,
normalization=self.normalization,
data_format=self.data_format,
activation=self.activation,
name="inters/%d" % i)
if -1 not in self.output_indices:
outputs = [
o for i, o in enumerate(outputs)
if i + 1 in self.output_indices
]
return tf.keras.Model(inputs=self.img_input,
outputs=outputs,
name=self.name)
def basic_block(
inputs,
filters,
strides=1,
dilation_rate=1,
data_format="channels_last",
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-3,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
trainable=True,
dropblock=None,
use_conv_shortcut=False,
expansion=1,
strides_in_1x1=False, # not use
name=None):
"""
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.
"""
x = ConvNormActBlock(filters=filters,
kernel_size=3,
strides=strides,
data_format=data_format,
dilation_rate=1 if strides > 1 else dilation_rate,
trainable=trainable,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/conv1")(inputs)
x = ConvNormActBlock(filters=filters,
kernel_size=3,
strides=1,
data_format=data_format,
dilation_rate=dilation_rate,
trainable=trainable,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/conv2")(x)
shortcut = inputs
if use_conv_shortcut:
shortcut = ConvNormActBlock(filters=filters,
kernel_size=1,
strides=strides,
data_format=data_format,
trainable=trainable,
normalization=normalization,
activation=None,
dropblock=dropblock,
name=name + "/shortcut")(shortcut)
x = tf.keras.layers.Add(name=name + "/add")([x, shortcut])
x = build_activation(**activation,
name=name + "/" + activation["activation"])(x)
return x
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 bottleneck(inputs,
filters,
in_filters,
radix=1,
cardinality=1,
bottleneck_width=64,
strides=1,
dilation_rate=1,
data_format="channels_last",
dropblock=None,
normalization=dict(normalization="batch_norm",
momentum=0.9,
epsilon=1e-5,
axis=-1,
trainable=True),
activation=dict(activation="relu"),
avg_down=False,
last_gamma=False,
avd=False,
avd_first=False,
is_first=False,
trainable=True,
name="bottleneck"):
avd = avd and (strides > 1 or is_first)
act = activation["activation"]
group_width = int(filters * (bottleneck_width / 64.)) * cardinality
x = ConvNormActBlock(filters=group_width,
kernel_size=(1, 1),
data_format=data_format,
trainable=trainable,
normalization=normalization,
activation=activation,
dropblock=dropblock,
name=name + "/conv1")(inputs)
if avd and avd_first:
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(x)
x = tf.keras.layers.AvgPool2D(pool_size=(3, 3),
strides=strides,
padding="valid",
data_format=data_format,
name=name + "/avd_layer")(x)
if radix > 1:
x = SplAtConv2d(x,
filters=group_width,
in_filters=group_width,
kernel_size=(3, 3),
strides=1 if avd else strides,
dilation_rate=dilation_rate,
cardinality=cardinality,
radix=radix,
reduction_factor=4,
data_format=data_format,
normalization=normalization,
dropblock=dropblock,
activation=activation,
trainable=trainable,
name=name + "/conv2")
else:
x = ConvNormActBlock(filters=group_width,
kernel_size=(3, 3),
strides=1 if avd else strides,
padding="same",
use_bias=False,
groups=cardinality,
trainable=trainable,
data_format=data_format,
normalization=normalization,
dropblock=dropblock,
activation=activation,
name=name + "/conv2")(x)
if avd and not avd_first:
x = tf.keras.layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(x)
x = tf.keras.layers.AvgPool2D(pool_size=(3, 3),
strides=strides,
padding="valid",
data_format=data_format,
name=name + "/avd_layer")(x)
x = ConvNormActBlock(filters=filters * 4,
kernel_size=(1, 1),
trainable=trainable,
data_format=data_format,
normalization=normalization,
dropblock=dropblock,
activation=None,
gamma_zeros=last_gamma,
name=name + "/conv3")(x)
shortcut = inputs
if strides != 1 or filters * 4 != in_filters:
if avg_down:
if dilation_rate == 1:
shortcut = tf.keras.layers.AvgPool2D(
pool_size=strides,
strides=strides,
padding="same",
data_format=data_format,
name=name + "/downsample/avgpool")(shortcut)
else:
shortcut = tf.keras.layers.AvgPool2D(
pool_size=1,
strides=1,
padding="same",
data_format=data_format,
name=name + "/downsample/avgpool")(shortcut)
strides = 1
shortcut = ConvNormActBlock(filters=filters * 4,
kernel_size=1,
strides=strides,
use_bias=False,
padding="same",
trainable=trainable,
data_format=data_format,
normalization=normalization,
dropblock=dropblock,
activation=None,
name=name + "/downsample")(shortcut)
x = tf.keras.layers.Add(name=name + "/sum")([x, shortcut])
x = build_activation(**activation, name=name + "/%s3" % act)(x)
return x
