def test_batch_norm(self, is_training, strategy):
inputs = tf.random.uniform([8, 40, 40, 3])
expect_results = utils.batch_norm_act(inputs, is_training, None)
# Call batch norm layer with is_training parameter.
bn_layer = utils_keras.build_batch_norm(is_training, strategy=strategy)
self.assertAllClose(expect_results, bn_layer(inputs, is_training))
def box_net(images,
level,
num_anchors,
num_filters,
is_training,
act_type,
repeats=4,
separable_conv=True,
survival_prob=None,
strategy=None,
data_format='channels_last'):
"""Box regression network."""
if separable_conv:
conv_op = functools.partial(
tf.layers.separable_conv2d, depth_multiplier=1,
data_format=data_format,
pointwise_initializer=tf.initializers.variance_scaling(),
depthwise_initializer=tf.initializers.variance_scaling())
else:
conv_op = functools.partial(
tf.layers.conv2d,
data_format=data_format,
kernel_initializer=tf.random_normal_initializer(stddev=0.01))
for i in range(repeats):
orig_images = images
images = conv_op(
images,
num_filters,
kernel_size=3,
activation=None,
bias_initializer=tf.zeros_initializer(),
padding='same',
name='box-%d' % i)
images = utils.batch_norm_act(
images,
is_training,
act_type=act_type,
init_zero=False,
strategy=strategy,
data_format=data_format,
name='box-%d-bn-%d' % (i, level))
if i > 0 and survival_prob:
images = utils.drop_connect(images, is_training, survival_prob)
images = images + orig_images
boxes = conv_op(
images,
4 * num_anchors,
kernel_size=3,
bias_initializer=tf.zeros_initializer(),
padding='same',
name='box-predict')
return boxes
def test_batchnorm_compatibility(self):
x = tf.Variable(tf.ones((4, 1, 1, 1)) * [[1.0], [2.0], [4.0], [8.0]])
for act_type in ['swish', 'swish_native', 'relu', 'relu6']:
bna = utils_keras.BatchNormAct(is_training_bn=False,
act_type=act_type)
self.assertEqual(
tf.reduce_sum(
utils.batch_norm_act(x,
is_training_bn=False,
act_type=act_type).numpy()),
tf.reduce_sum(bna.call(x).numpy()))
def _maybe_apply_1x1(feat):
"""Apply 1x1 conv to change layer width if necessary."""
if num_channels != target_num_channels:
feat = tf.keras.layers.Conv2D(feat,
filters=target_num_channels,
kernel_size=(1, 1),
padding='same')
feat = utils.batch_norm_act(feat,
is_training_bn=is_training,
act_type=None,
name='bn')
return feat
def _maybe_apply_1x1(self, feat):
"""Apply 1x1 conv to change layer width if necessary."""
if self.num_channels != self.target_num_channels:
feat = self.conv2d(feat)
if self.apply_bn:
feat = utils.batch_norm_act(feat,
is_training_bn=self.is_training,
act_type=None,
data_format=self.data_format,
use_tpu=self.use_tpu,
name='bn')
return feat
def class_net(images,
level,
num_classes,
num_anchors,
num_filters,
is_training,
act_type,
separable_conv=True,
repeats=4,
survival_prob=None,
use_tpu=False):
"""Class prediction network."""
if separable_conv:
conv_op = functools.partial(
tf.layers.separable_conv2d,
depth_multiplier=1,
pointwise_initializer=tf.initializers.variance_scaling(),
depthwise_initializer=tf.initializers.variance_scaling())
else:
conv_op = functools.partial(
tf.layers.conv2d,
kernel_initializer=tf.random_normal_initializer(stddev=0.01))
for i in range(repeats):
orig_images = images
images = conv_op(images,
num_filters,
kernel_size=3,
bias_initializer=tf.zeros_initializer(),
activation=None,
padding='same',
name='class-%d' % i)
images = utils.batch_norm_act(images,
is_training,
act_type=act_type,
init_zero=False,
use_tpu=use_tpu,
name='class-%d-bn-%d' % (i, level))
if i > 0 and survival_prob:
images = utils.drop_connect(images, is_training, survival_prob)
images = images + orig_images
classes = conv_op(
images,
num_classes * num_anchors,
kernel_size=3,
bias_initializer=tf.constant_initializer(-np.log((1 - 0.01) / 0.01)),
padding='same',
name='class-predict')
return classes
def _maybe_apply_1x1(feat):
"""Apply 1x1 conv to change layer width if necessary."""
if num_channels != target_num_channels:
feat = tf.layers.conv2d(feat,
filters=target_num_channels,
kernel_size=(1, 1),
padding='same')
if apply_bn:
feat = utils.batch_norm_act(feat,
is_training_bn=is_training,
act_type=None,
data_format='channels_last',
use_tpu=use_tpu,
name='bn')
return feat
def box_net(images,
level,
num_anchors,
num_filters,
is_training,
act_type,
repeats=4,
separable_conv=True,
survival_prob=None):
"""Box regression network."""
for i in range(repeats):
orig_images = images
images = tf.keras.layers.SeparableConv2D(
images,
num_filters,
kernel_size=3,
activation=None,
bias_initializer=tf.zeros_initializer(),
padding='same',
name='box-%d' % i,
depth_multiplier=1,
pointwise_initializer=tf.keras.initializers.VarianceScaling(),
depthwise_initializer=tf.keras.initializers.VarianceScaling())
images = utils.batch_norm_act(images,
is_training,
act_type=act_type,
init_zero=False,
name='box-%d-bn-%d' % (i, level))
if i > 0 and survival_prob:
images = utils.drop_connect(images, is_training, survival_prob)
images = images + orig_images
boxes = tf.keras.layers.SeparableConv2D(
images,
4 * num_anchors,
kernel_size=3,
bias_initializer=tf.zeros_initializer(),
padding='same',
name='box-predict',
depth_multiplier=1,
pointwise_initializer=tf.keras.initializers.VarianceScaling(),
depthwise_initializer=tf.keras.initializers.VarianceScaling())
return boxes
def class_net(images,
level,
num_classes,
num_anchors,
num_filters,
is_training,
act_type,
repeats=4,
survival_prob=None):
"""Class prediction network."""
for i in range(repeats):
orig_images = images
images = tf.keras.layers.SeparableConv2D(
images,
num_filters,
kernel_size=3,
bias_initializer=tf.keras.initializers.zeros(),
activation=None,
padding='same',
name='class-%d' % i,
depth_multiplier=1,
pointwise_initializer=tf.keras.initializers.VarianceScaling(),
depthwise_initializer=tf.keras.initializers.VarianceScaling())
images = utils.batch_norm_act(images,
is_training,
act_type=act_type,
init_zero=False,
name='class-%d-bn-%d' % (i, level))
if i > 0 and survival_prob:
images = utils.drop_connect(images, is_training, survival_prob)
images = images + orig_images
classes = tf.keras.layers.SeparableConv2D(
images,
num_classes * num_anchors,
kernel_size=3,
bias_initializer=tf.keras.initializers.Constant(-np.log((1 - 0.01) /
0.01)),
padding='same',
name='class-predict',
depth_multiplier=1,
pointwise_initializer=tf.keras.initializers.VarianceScaling(),
depthwise_initializer=tf.keras.initializers.VarianceScaling())
return classes
def build_bifpn_layer(feats, feat_sizes, config):
"""Builds a feature pyramid given previous feature pyramid and config."""
p = config # use p to denote the network config.
if p.fpn_config:
fpn_config = p.fpn_config
else:
fpn_config = fpn_configs.get_fpn_config(p.fpn_name, p.min_level,
p.max_level,
p.fpn_weight_method)
num_output_connections = [0 for _ in feats]
for i, fnode in enumerate(fpn_config.nodes):
with tf.variable_scope('fnode{}'.format(i)):
logging.info('fnode %d : %s', i, fnode)
new_node_height = feat_sizes[fnode['feat_level']]['height']
new_node_width = feat_sizes[fnode['feat_level']]['width']
nodes = []
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,
'{}_{}_{}'.format(idx, input_offset, len(feats)),
new_node_height,
new_node_width,
p.fpn_num_filters,
p.apply_bn_for_resampling,
p.is_training_bn,
p.conv_after_downsample,
strategy=p.strategy,
data_format=config.data_format)
nodes.append(input_node)
new_node = fuse_features(nodes, fpn_config.weight_method)
with tf.variable_scope('op_after_combine{}'.format(len(feats))):
if not p.conv_bn_act_pattern:
new_node = utils.activation_fn(new_node, p.act_type)
if p.separable_conv:
conv_op = functools.partial(tf.layers.separable_conv2d,
depth_multiplier=1)
else:
conv_op = tf.layers.conv2d
new_node = conv_op(new_node,
filters=p.fpn_num_filters,
kernel_size=(3, 3),
padding='same',
use_bias=not p.conv_bn_act_pattern,
data_format=config.data_format,
name='conv')
new_node = utils.batch_norm_act(
new_node,
is_training_bn=p.is_training_bn,
act_type=None if not p.conv_bn_act_pattern else p.act_type,
data_format=config.data_format,
strategy=p.strategy,
name='bn')
feats.append(new_node)
num_output_connections.append(0)
output_feats = {}
for l in range(p.min_level, p.max_level + 1):
for i, fnode in enumerate(reversed(fpn_config.nodes)):
if fnode['feat_level'] == l:
output_feats[l] = feats[-1 - i]
break
return output_feats
def build_bifpn_layer(feats, feat_sizes, config):
"""Builds a feature pyramid given previous feature pyramid and config."""
p = config # use p to denote the network config.
if p.fpn_config:
fpn_config = p.fpn_config
else:
fpn_config = get_fpn_config(p.fpn_name, p.min_level, p.max_level,
p.fpn_weight_method)
num_output_connections = [0 for _ in feats]
for i, fnode in enumerate(fpn_config.nodes):
with tf.variable_scope('fnode{}'.format(i)):
logging.info('fnode %d : %s', i, fnode)
new_node_height = feat_sizes[fnode['feat_level']]['height']
new_node_width = feat_sizes[fnode['feat_level']]['width']
nodes = []
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, '{}_{}_{}'.format(idx, input_offset, len(feats)),
new_node_height, new_node_width, p.fpn_num_filters,
p.apply_bn_for_resampling, p.is_training_bn,
p.conv_after_downsample,
p.use_native_resize_op,
p.pooling_type,
data_format=config.data_format)
nodes.append(input_node)
# Combine all nodes.
dtype = nodes[0].dtype
if fpn_config.weight_method == 'attn':
edge_weights = [tf.cast(tf.Variable(1.0, name='WSM'), dtype=dtype)
for _ in range(len(fnode['inputs_offsets']))]
normalized_weights = tf.nn.softmax(tf.stack(edge_weights))
nodes = tf.stack(nodes, axis=-1)
new_node = tf.reduce_sum(tf.multiply(nodes, normalized_weights), -1)
elif fpn_config.weight_method == 'fastattn':
edge_weights = [
tf.nn.relu(tf.cast(tf.Variable(1.0, name='WSM'), dtype=dtype))
for _ in range(len(fnode['inputs_offsets']))
]
weights_sum = tf.add_n(edge_weights)
nodes = [nodes[i] * edge_weights[i] / (weights_sum + 0.0001)
for i in range(len(nodes))]
new_node = tf.add_n(nodes)
elif fpn_config.weight_method == 'sum':
new_node = tf.add_n(nodes)
else:
raise ValueError(
'unknown weight_method {}'.format(fpn_config.weight_method))
with tf.variable_scope('op_after_combine{}'.format(len(feats))):
if not p.conv_bn_act_pattern:
new_node = utils.activation_fn(new_node, p.act_type)
if p.separable_conv:
conv_op = functools.partial(
tf.layers.separable_conv2d, depth_multiplier=1)
else:
conv_op = tf.layers.conv2d
new_node = conv_op(
new_node,
filters=p.fpn_num_filters,
kernel_size=(3, 3),
padding='same',
use_bias=True if not p.conv_bn_act_pattern else False,
data_format=config.data_format,
name='conv')
new_node = utils.batch_norm_act(
new_node,
is_training_bn=p.is_training_bn,
act_type=None if not p.conv_bn_act_pattern else p.act_type,
data_format=config.data_format,
use_tpu=p.use_tpu,
name='bn')
feats.append(new_node)
num_output_connections.append(0)
output_feats = {}
for l in range(p.min_level, p.max_level + 1):
for i, fnode in enumerate(reversed(fpn_config.nodes)):
if fnode['feat_level'] == l:
output_feats[l] = feats[-1 - i]
break
return output_feats
