def residual_conv(prev,
level,
pad=1,
lvl=1,
sub_lvl=1,
modify_stride=False,
is_training=True,
use_gn=False):
lvl = str(lvl)
sub_lvl = str(sub_lvl)
names = [
"conv" + lvl + "_" + sub_lvl + "_1x1_reduce",
"conv" + lvl + "_" + sub_lvl + "_1x1_reduce_bn",
"conv" + lvl + "_" + sub_lvl + "_3x3",
"conv" + lvl + "_" + sub_lvl + "_3x3_bn",
"conv" + lvl + "_" + sub_lvl + "_1x1_increase",
"conv" + lvl + "_" + sub_lvl + "_1x1_increase_bn"
]
if modify_stride is False:
prev = Conv2D(prev,
64 * level, (1, 1),
strides=(1, 1),
name=names[0],
use_bias=False)
elif modify_stride is True:
prev = Conv2D(prev,
64 * level, (1, 1),
strides=(2, 2),
name=names[0],
use_bias=False)
prev = GroupNorm(prev, 4 * level, name=names[0] + '_gn') if use_gn else BN(
prev, name=names[1], is_training=is_training)
prev = Activation(prev, 'relu')
prev = ZeroPadding2D_symmetric(prev, padding=pad)
prev = Conv2D(prev,
64 * level, (3, 3),
strides=(1, 1),
dilation_rate=pad,
name=names[2],
use_bias=False)
prev = GroupNorm(prev, 4 * level, name=names[2] + '_gn') if use_gn else BN(
prev, name=names[3], is_training=is_training)
prev = Activation(prev, 'relu')
prev = Conv2D(prev,
256 * level, (1, 1),
strides=(1, 1),
name=names[4],
use_bias=False)
prev = GroupNorm(prev, 16 * level, name=names[4] +
'_gn') if use_gn else BN(
prev, name=names[5], is_training=is_training)
return prev
def short_convolution_branch(prev,
level,
lvl=1,
sub_lvl=1,
modify_stride=False,
is_training=True,
use_gn=False):
lvl = str(lvl)
sub_lvl = str(sub_lvl)
names = [
"conv" + lvl + "_" + sub_lvl + "_1x1_proj",
"conv" + lvl + "_" + sub_lvl + "_1x1_proj_bn"
]
if modify_stride is False:
prev = Conv2D(prev,
256 * level, (1, 1),
strides=(1, 1),
name=names[0],
use_bias=False)
elif modify_stride is True:
prev = Conv2D(prev,
256 * level, (1, 1),
strides=(2, 2),
name=names[0],
use_bias=False)
prev = GroupNorm(prev, 16 * level, name=names[0] +
'_gn') if use_gn else BN(
prev, name=names[1], is_training=is_training)
return prev
def build_pspnet_prior(nb_classes,
inputs,
activation='softmax',
is_training=True):
"""Build PSPNet."""
dropout_ratio = 0.1 if is_training else 0
x = Conv2D(inputs,
512, (3, 3),
strides=(1, 1),
padding="same",
name="conv5_4",
use_bias=False)
x = BN(x, name="conv5_4_bn", is_training=is_training)
x = Activation(x, 'relu')
x = Dropout(x, dropout_ratio)
x = Conv2D(x,
nb_classes, (1, 1),
strides=(1, 1),
padding="same",
name="conv6")
prior = Activation(x, activation)
return Interp(prior, [713, 713])
def build_pspnet50_prior_on_resnet_stem(nb_classes,
inputs,
activation='softmax',
is_training=True):
"""Build PSPNet."""
dropout_ratio = 0.1 if is_training else 0
inp = inputs
res = residual_short(inp,
8,
pad=4,
lvl=5,
sub_lvl=1,
is_training=is_training)
for i in range(2):
res = residual_empty(res,
8,
pad=4,
lvl=5,
sub_lvl=i + 2,
is_training=is_training)
res = Activation(res, 'relu')
psp = build_pyramid_pooling_module(res, is_training=is_training)
x = Conv2D(psp,
512, (3, 3),
strides=(1, 1),
padding="same",
name="conv5_4",
use_bias=False)
x = BN(x, name="conv5_4_bn", is_training=is_training)
x = Activation(x, 'relu')
x = Dropout(x, dropout_ratio)
x = Conv2D(x,
nb_classes, (1, 1),
strides=(1, 1),
padding="same",
name="conv6")
prior = Activation(x, activation)
return Interp(prior, [473, 473])
def build_pspnet_ade20k_prior(nb_classes,
inputs,
activation='softmax',
is_training=True):
"""Build PSPNet."""
dropout_ratio = 0.1 if is_training else 0
x_dd = Conv2D(inputs,
512, (3, 3),
strides=(1, 1),
padding="same",
name="conv5_4",
use_bias=False)
x_dd = BN(x_dd, name="conv5_4_bn", is_training=is_training)
x_dd = Activation(x_dd, 'relu')
x_dd = Dropout(x_dd, dropout_ratio)
prior = Conv2D(x_dd,
nb_classes, (1, 1),
strides=(1, 1),
padding="same",
name="conv6")
prior = Activation(prior, activation)
x_nd = Conv2D(inputs,
512, (3, 3),
strides=(1, 1),
padding="same",
name="base_prediction_conv5_4",
use_bias=False)
x_nd = GroupNorm(x_nd, 32, name="base_prediction_conv5_4_gn")
x_nd = Activation(x_nd, 'relu')
x_nd = Dropout(x_nd, dropout_ratio)
base_prediction = Conv2D(x_nd,
nb_classes, (1, 1),
strides=(1, 1),
padding="same",
name="base_prediction_conv6")
return base_prediction, Interp(prior, [473, 473])
def interp_block(prev_layer,
level,
feature_map_shape,
str_lvl=1,
dims=None,
name=None,
is_training=True,
use_gn=False):
if name:
names = [name + "_conv", name + "_conv_bn"]
else:
str_lvl = str(str_lvl)
names = [
"conv5_3_pool" + str_lvl + "_conv",
"conv5_3_pool" + str_lvl + "_conv_bn"
]
out_dims = dims if dims else 512
kernel = (int(ceil(feature_map_shape[0] / level)),
int(ceil(feature_map_shape[1] / level)))
strides = (int(ceil(feature_map_shape[0] / level)),
int(ceil(feature_map_shape[1] / level)))
prev_layer = AveragePooling2D(prev_layer, kernel, strides=strides)
prev_layer = Conv2D(prev_layer,
out_dims, (1, 1),
strides=(1, 1),
name=names[0],
use_bias=False)
prev_layer = GroupNorm(
prev_layer, out_dims / 16, name=names[0] + '_gn') if use_gn else BN(
prev_layer, name=names[1], is_training=is_training)
prev_layer = Activation(prev_layer, 'relu')
prev_layer = Interp(prev_layer, shape=feature_map_shape)
return prev_layer
def build_pspnet(resnet_layers,
inputs,
nb_classes,
dynamic_envs,
activation='softmax',
is_training=True):
dropout_ratio = 0.1 if is_training else 0
inp = tf.concat([inputs, dynamic_envs], axis=-1)
inputs_shape_spatial = inp.shape.as_list()[1:-1]
res = ResNet(inp, layers=resnet_layers, is_training=is_training)
psp = build_pyramid_pooling_module(res, is_training=is_training)
'''
# Prior
x_prior = Conv2D(psp, 512, (3, 3), strides=(1, 1), padding="same", name="conv5_4",
use_bias=False)
x_prior = BN(x_prior, name="conv5_4_bn", is_training=is_training)
x_prior = Activation(x_prior, 'relu')
x_prior = Dropout(x_prior, dropout_ratio)
x_prior = Conv2D(x_prior, nb_classes, (1, 1), strides=(1, 1), padding="same", name="conv6")
prior = Activation(x_prior, activation)
'''
# Top
with tf.variable_scope('top') as _:
x = Conv2D(psp,
1024, (3, 3),
strides=(1, 1),
padding="same",
name="conv_top_1",
use_bias=False)
x = BN(x, name="conv_top_1_bn", is_training=is_training)
x = Activation(x, 'relu')
top_shortcut = x
x = Conv2D(x,
1024, (3, 3),
strides=(1, 1),
padding="same",
name="conv_top_2",
use_bias=False)
x = BN(x, name="conv_top_2_bn", is_training=is_training)
x = Activation(x, 'relu')
x = Conv2D(x,
1024, (3, 3),
strides=(1, 1),
padding="same",
name="conv_top_3",
use_bias=False)
x = BN(x, name="conv_top_3_bn", is_training=is_training)
x = Activation(x, 'relu')
x = x + top_shortcut
x = Dropout(x, dropout_ratio)
preact_policy = Conv2D(x,
nb_classes, (1, 1),
strides=(1, 1),
padding="same",
name="conv_class")
interp_preact_policy = Interp(x=preact_policy,
shape=inputs_shape_spatial)
state_value = Conv2D(x,
1, (1, 1),
strides=(1, 1),
name="conv_value",
use_bias=False)
policy = Activation(preact_policy, 'softmax')
interp_policy = Activation(interp_preact_policy, 'softmax')
return tf.squeeze(state_value,
axis=-1), (preact_policy, policy, interp_preact_policy,
interp_policy)
def ResNet_stem(inputs, is_training=True):
# Names for the first couple layers of model_util
names = [
"conv1_1_3x3_s2", "conv1_1_3x3_s2_bn", "conv1_2_3x3", "conv1_2_3x3_bn",
"conv1_3_3x3", "conv1_3_3x3_bn"
]
# Short branch(only start of network)
inp = inputs
cnv1 = Conv2D(inp,
64, (3, 3),
strides=(2, 2),
padding='same',
name=names[0],
use_bias=False) # "conv1_1_3x3_s2"
bn1 = BN(cnv1, name=names[1],
is_training=is_training) # "conv1_1_3x3_s2/bn"
relu1 = Activation(bn1, 'relu') # "conv1_1_3x3_s2/relu"
cnv1 = Conv2D(relu1,
64, (3, 3),
strides=(1, 1),
padding='same',
name=names[2],
use_bias=False) # "conv1_2_3x3"
bn1 = BN(cnv1, name=names[3], is_training=is_training) # "conv1_2_3x3/bn"
relu1 = Activation(bn1, 'relu') # "conv1_2_3x3/relu"
cnv1 = Conv2D(relu1,
128, (3, 3),
strides=(1, 1),
padding='same',
name=names[4],
use_bias=False) # "conv1_3_3x3"
bn1 = BN(cnv1, name=names[5], is_training=is_training) # "conv1_3_3x3/bn"
relu1 = Activation(bn1, 'relu') # "conv1_3_3x3/relu"
res = MaxPooling2D(relu1, pool_size=(3, 3), padding='same',
strides=(2, 2)) # "pool1_3x3_s2"
# ---Residual layers(body of network)
"""
Modify_stride --Used only once in first 3_1 convolutions block.
changes stride of first convolution from 1 -> 2
"""
# 2_1- 2_3
res = residual_short(res,
1,
pad=1,
lvl=2,
sub_lvl=1,
is_training=is_training)
for i in range(2):
res = residual_empty(res,
1,
pad=1,
lvl=2,
sub_lvl=i + 2,
is_training=is_training)
# 3_1 - 3_3
res = residual_short(res,
2,
pad=1,
lvl=3,
sub_lvl=1,
modify_stride=True,
is_training=is_training)
for i in range(3):
res = residual_empty(res,
2,
pad=1,
lvl=3,
sub_lvl=i + 2,
is_training=is_training)
# 4_1 - 4_6
res = residual_short(res,
4,
pad=2,
lvl=4,
sub_lvl=1,
is_training=is_training)
for i in range(5):
res = residual_empty(res,
4,
pad=2,
lvl=4,
sub_lvl=i + 2,
is_training=is_training)
return res