def call(self, x, training=None):
if self.downsample:
y1 = self.dw_conv4(x)
y1 = self.dw_bn4(y1, training=training)
y1 = self.expand_conv5(y1)
y1 = self.expand_bn5(y1, training=training)
y1 = self.activ(y1)
x2 = x
else:
y1, x2 = tf.split(x,
num_or_size_splits=2,
axis=get_channel_axis(self.data_format))
y2 = self.compress_conv1(x2)
y2 = self.compress_bn1(y2, training=training)
y2 = self.activ(y2)
y2 = self.dw_conv2(y2)
y2 = self.dw_bn2(y2, training=training)
y2 = self.expand_conv3(y2)
y2 = self.expand_bn3(y2, training=training)
y2 = self.activ(y2)
if self.use_se:
y2 = self.se(y2)
if self.use_residual and not self.downsample:
y2 = y2 + x2
x = tf.concat([y1, y2], axis=get_channel_axis(self.data_format))
x = self.c_shuffle(x)
return x
def call(self, x, x0=None, training=None):
x1 = self.branch1(x, training=training)
x2 = self.branch2(x1, training=training)
x3 = self.branch3(x2, training=training)
if self.resize:
y = tf.concat([x1, x2, x3], axis=get_channel_axis(self.data_format))
y = self.last_conv(y, training=training)
return y, y
else:
y = tf.concat([x1, x2, x3, x0], axis=get_channel_axis(self.data_format))
y = self.last_conv(y, training=training)
return y, x0
def call(self, x, training=None):
y = self.body(x, training=training)
heatmap = self.heatmap_bend(y, training=training)
paf = self.paf_bend(y, training=training)
y = tf.concat((x, heatmap, paf),
axis=get_channel_axis(self.data_format))
return y
def call(self, x, training=None):
identity = x
# pointwise group convolution 1
x = self.compress_conv1(x)
x = self.compress_bn1(x, training=training)
x = self.activ(x)
x = self.c_shuffle(x)
# merging
y = self.s_merge_conv(x)
y = self.s_merge_bn(y, training=training)
y = self.activ(y)
# depthwise convolution (bottleneck)
x = self.dw_conv2(x)
x = self.dw_bn2(x, training=training)
# evolution
y = self.s_conv(y)
y = self.s_conv_bn(y, training=training)
y = self.activ(y)
y = self.s_evolve_conv(y)
y = self.s_evolve_bn(y, training=training)
y = tf.nn.sigmoid(y)
x = x * y
# pointwise group convolution 2
x = self.expand_conv3(x)
x = self.expand_bn3(x, training=training)
# identity branch
if self.downsample:
identity = self.avgpool(identity)
x = tf.concat([x, identity],
axis=get_channel_axis(self.data_format))
else:
x = x + identity
x = self.activ(x)
return x
def __init__(self,
in_channels,
out_channels,
kernel_size,
strides,
padding,
dilation=1,
groups=1,
use_bias=False,
use_bn=True,
bn_eps=1e-5,
activation=(lambda: nn.ReLU()),
data_format="channels_last",
**kwargs):
super(MixConvBlock, self).__init__(**kwargs)
self.activate = (activation is not None)
self.use_bn = use_bn
self.conv = MixConv(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
strides=strides,
padding=padding,
dilation=dilation,
groups=groups,
use_bias=use_bias,
axis=get_channel_axis(data_format),
data_format=data_format,
name="conv")
if self.use_bn:
self.bn = BatchNorm(epsilon=bn_eps,
data_format=data_format,
name="bn")
if self.activate:
self.activ = get_activation_layer(activation)
def call(self, x, training=None):
identity = x
x = self.conv(x, training=training)
if self.use_dropout:
x = self.dropout(x, training=training)
x = tf.concat([identity, x], axis=get_channel_axis(self.data_format))
return x
def call(self, x1, x2=None, training=None):
axis = get_channel_axis(self.data_format)
x_in = tf.concat([x1, x2], axis=axis) if x2 is not None else x1
if self.has_proj:
x_s = self.conv_proj(x_in, training=training)
channels = (x_s.get_shape().as_list())[axis]
x_s1, x_s2 = tf.split(
x_s,
num_or_size_splits=[self.bw, channels - self.bw],
axis=axis)
# x_s1 = F.slice_axis(x_s, axis=1, begin=0, end=self.bw)
# x_s2 = F.slice_axis(x_s, axis=1, begin=self.bw, end=None)
else:
assert (x2 is not None)
x_s1 = x1
x_s2 = x2
x_in = self.conv1(x_in, training=training)
x_in = self.conv2(x_in, training=training)
if self.b_case:
x_in = self.preactiv(x_in, training=training)
y1 = self.conv3a(x_in, training=training)
y2 = self.conv3b(x_in, training=training)
else:
x_in = self.conv3(x_in, training=training)
# y1 = F.slice_axis(x_in, axis=1, begin=0, end=self.bw)
# y2 = F.slice_axis(x_in, axis=1, begin=self.bw, end=None)
channels = (x_in.get_shape().as_list())[axis]
y1, y2 = tf.split(x_in,
num_or_size_splits=[self.bw, channels - self.bw],
axis=axis)
residual = x_s1 + y1
dense = tf.concat([x_s2, y2], axis=axis)
return residual, dense
def call(self, x, training=None):
axis = get_channel_axis(self.data_format)
layer_outs = [x]
for links_i, layer_i in zip(self.links_list, self.blocks.children):
layer_in = []
for idx_ij in links_i:
layer_in.append(layer_outs[idx_ij])
if len(layer_in) > 1:
x = tf.concat(layer_in, axis=axis)
else:
x = layer_in[0]
out = layer_i(x, training=training)
layer_outs.append(out)
outs = []
for i, layer_out_i in enumerate(layer_outs):
if (i == len(layer_outs) - 1) or (i % 2 == 1):
outs.append(layer_out_i)
x = tf.concat(outs, axis=axis)
if self.use_dropout:
x = self.dropout(x, training=training)
x = self.conv(x, training=training)
if self.downsampling:
x = self.downsample(x, training=training)
return x
def call(self, x, training=None):
x = self.conv1(x, training=training)
x = self.res1(x, training=training)
x = self.pool(x, training=training)
x = self.res2(x, training=training)
y = self.dilation_branch(x, training=training)
x = tf.concat([x, y], axis=get_channel_axis(self.data_format))
return x
def __init__(self, data_format="channels_last", **kwargs):
super(GlobalAvgMaxPool2D, self).__init__(**kwargs)
self.axis = get_channel_axis(data_format)
self.avg_pool = nn.GlobalAvgPool2D(data_format=data_format,
name="avg_pool")
self.max_pool = nn.GlobalMaxPool2D(data_format=data_format,
name="max_pool")
def call(self, x, x0, training=None):
y1 = self.pool(x)
y2, _ = self.eesp(x, None, training=training)
x = tf.concat([y1, y2], axis=get_channel_axis(self.data_format))
x0 = self.pool(x0)
y3 = self.shortcut_block(x0, training=training)
x = x + y3
x = self.activ(x)
return x, x0
def call(self, x, training=None):
axis = get_channel_axis(self.data_format)
att1 = tf.math.reduce_max(x, axis=axis, keepdims=True)
att2 = tf.math.reduce_mean(x, axis=axis, keepdims=True)
att = tf.concat([att1, att2], axis=axis)
att = self.conv(att, training=training)
att = tf.broadcast_to(self.sigmoid(att), shape=x.shape)
x = x * att
return x
def call(self, x2, x1, extra, training=None):
last_branch = x2
x = tf.concat([x2, x1] + list(extra),
axis=get_channel_axis(self.data_format))
x = self.conv(x, training=training)
if self.residual:
x += last_branch
x = self.activ(x)
return x
def call(self, x):
if self.residual:
identity = x
x = self.squeeze(x)
y1 = self.expand1x1(x)
y2 = self.expand3x3(x)
out = tf.concat([y1, y2], axis=get_channel_axis(self.data_format))
if self.residual:
out = out + identity
return out
def call(self, x, training=None):
if self.downsample:
y1 = self.shortcut_dconv(x, training=training)
y1 = self.shortcut_conv(y1, training=training)
x2 = x
else:
y1, x2 = tf.split(x,
num_or_size_splits=2,
axis=get_channel_axis(self.data_format))
y2 = self.conv1(x2, training=training)
y2 = self.dconv(y2, training=training)
y2 = self.conv2(y2, training=training)
if self.use_se:
y2 = self.se(y2)
if self.use_residual and not self.downsample:
y2 = y2 + x2
x = tf.concat([y1, y2], axis=get_channel_axis(self.data_format))
x = self.c_shuffle(x)
return x
def call(self, x, training=None):
if is_channels_first(self.data_format):
features = x[:, :self.features_channels]
else:
features = x[:, :, :, :self.features_channels]
y = self.body(x, training=training)
heatmap = self.heatmap_bend(y, training=training)
paf = self.paf_bend(y, training=training)
y = tf.concat((features, heatmap, paf),
axis=get_channel_axis(self.data_format))
return y
def call(self, x, training=None):
if is_channels_first(self.data_format):
heatmap_paf_2d = x[:, -self.num_heatmap_paf:]
else:
heatmap_paf_2d = x[:, :, :, -self.num_heatmap_paf:]
if not self.calc_3d_features:
return heatmap_paf_2d
x = self.body(x, training=training)
x = self.features_bend(x, training=training)
y = tf.concat((heatmap_paf_2d, x),
axis=get_channel_axis(self.data_format))
return y
def call(self, x, y, training=None):
x = self.up(x)
x = self.bn(x, training=None)
w_conf = tf.nn.softmax(x)
axis = get_channel_axis(self.data_format)
# w_max = tf.broadcast_to(tf.expand_dims(tf.reduce_max(w_conf, axis=axis), axis=axis), shape=x.shape)
w_max = tf.repeat(tf.expand_dims(tf.reduce_max(w_conf, axis=axis),
axis=axis),
x.shape[axis],
axis=axis)
x = y * (1 - w_max) + x
return x
def call(self, x, training=None):
identity = x
x = self.compress_conv1(x)
x = self.compress_bn1(x, training=training)
x = self.activ(x)
x = self.c_shuffle(x)
x = self.dw_conv2(x)
x = self.dw_bn2(x, training=training)
x = self.expand_conv3(x)
x = self.expand_bn3(x, training=training)
if self.downsample:
identity = self.avgpool(identity)
x = tf.concat([x, identity],
axis=get_channel_axis(self.data_format))
else:
x = x + identity
x = self.activ(x)
return x
def cell_forward(self, x, x_prev, training=None):
assert (hasattr(self, 'comb0_left'))
x_left = x_prev
x_right = x
x0 = self.comb0_left(x_left, training=training) + self.comb0_right(
x_left, training=training)
x1 = self.comb1_left(x_right, training=training) + self.comb1_right(
x_right, training=training)
x2 = self.comb2_left(x_right, training=training) + self.comb2_right(
x_right, training=training)
x3 = self.comb3_left(x2, training=training) + self.comb3_right(
x_right, training=training)
x4 = self.comb4_left(x_left, training=training) + (self.comb4_right(
x_right, training=training) if self.comb4_right else x_right)
x_out = tf.concat([x0, x1, x2, x3, x4],
axis=get_channel_axis(self.data_format))
return x_out
def call(self, x, training=None):
x = self.conv_list(x, training=training)
y1 = self.conv1x3(x, training=training)
y2 = self.conv3x1(x, training=training)
x = tf.concat([y1, y2], axis=get_channel_axis(self.data_format))
return x
def call(self, x, training=None):
x1 = self.branch1(x, training=training)
x2 = self.branch2(x, training=training)
x = tf.concat([x, x1, x2], axis=get_channel_axis(self.data_format))
return x
def call(self, x, training=None):
x = self.down_conv(x, training=None)
y = self.main_branch(x, training=None)
x = tf.concat([x, y], axis=get_channel_axis(self.data_format))
x = self.preactiv(x, training=None)
return x, y
def __init__(self,
data_format="channels_last",
**kwargs):
super(HierarchicalConcurrent, self).__init__(**kwargs)
self.axis = get_channel_axis(data_format)
def call(self, x1, x2, training=None):
assert (x2 is not None)
x = tf.concat([x1, x2], axis=get_channel_axis(self.data_format))
x = self.activ(x)
return x, None
def call(self, x, training=None):
x = self.main_conv(x, training=training)
y = self.cheap_conv(x, training=training)
return tf.concat([x, y], axis=get_channel_axis(self.data_format))
