# vim: expandtab:ts=4:sw=4
import argparse
import tensorflow as tf
import tensorflow.contrib.slim as slim
def _batch_norm_fn(x, scope=None):
if scope is None:
scope = tf.get_variable_scope().name + "/bn"
return slim.batch_norm(x, scope=scope)
def create_link(
incoming, network_builder, scope, nonlinearity=tf.nn.elu,
weights_initializer=tf.truncated_normal_initializer(stddev=1e-3),
regularizer=None, is_first=False, summarize_activations=True):
if is_first:
network = incoming
else:
network = _batch_norm_fn(incoming, scope=scope + "/bn")
network = nonlinearity(network)
if summarize_activations:
tf.summary.histogram(scope+"/activations", network)
pre_block_network = network
post_block_network = network_builder(pre_block_network, scope)
incoming_dim = pre_block_network.get_shape().as_list()[-1]
outgoing_dim = post_block_network.get_shape().as_list()[-1]
if incoming_dim != outgoing_dim:
assert outgoing_dim == 2 * incoming_dim, \
"%d != %d" % (outgoing_dim, 2 * incoming)
projection = slim.conv2d(
incoming, outgoing_dim, 1, 2, padding="SAME", activation_fn=None,
scope=scope+"/projection", weights_initializer=weights_initializer,
biases_initializer=None, weights_regularizer=regularizer)
network = projection + post_block_network
else:
network = incoming + post_block_network
return network
def create_inner_block(
incoming, scope, nonlinearity=tf.nn.elu,
weights_initializer=tf.truncated_normal_initializer(1e-3),
bias_initializer=tf.zeros_initializer(), regularizer=None,
increase_dim=False, summarize_activations=True):
n = incoming.get_shape().as_list()[-1]
stride = 1
if increase_dim:
n *= 2
stride = 2
incoming = slim.conv2d(
incoming, n, [3, 3], stride, activation_fn=nonlinearity, padding="SAME",
normalizer_fn=_batch_norm_fn, weights_initializer=weights_initializer,
biases_initializer=bias_initializer, weights_regularizer=regularizer,
scope=scope + "/1")
if summarize_activations:
tf.summary.histogram(incoming.name + "/activations", incoming)
incoming = slim.dropout(incoming, keep_prob=0.6)
incoming = slim.conv2d(
incoming, n, [3, 3], 1, activation_fn=None, padding="SAME",
normalizer_fn=None, weights_initializer=weights_initializer,
biases_initializer=bias_initializer, weights_regularizer=regularizer,
scope=scope + "/2")
return incoming
def residual_block(incoming, scope, nonlinearity=tf.nn.elu,
weights_initializer=tf.truncated_normal_initializer(1e3),
bias_initializer=tf.zeros_initializer(), regularizer=None,
increase_dim=False, is_first=False,
summarize_activations=True):
def network_builder(x, s):
return create_inner_block(
x, s, nonlinearity, weights_initializer, bias_initializer,
regularizer, increase_dim, summarize_activations)
return create_link(
incoming, network_builder, scope, nonlinearity, weights_initializer,
regularizer, is_first, summarize_activations)
def _create_network(incoming, reuse=None, weight_decay=1e-8):
nonlinearity = tf.nn.elu
conv_weight_init = tf.truncated_normal_initializer(stddev=1e-3)
conv_bias_init = tf.zeros_initializer()
conv_regularizer = slim.l2_regularizer(weight_decay)
fc_weight_init = tf.truncated_normal_initializer(stddev=1e-3)
fc_bias_init = tf.zeros_initializer()
fc_regularizer = slim.l2_regularizer(weight_decay)
def batch_norm_fn(x):
return slim.batch_norm(x, scope=tf.get_variable_scope().name + "/bn")
network = incoming
network = slim.conv2d(
network, 32, [3, 3], stride=1, activation_fn=nonlinearity,
padding="SAME", normalizer_fn=batch_norm_fn, scope="conv1_1",
weights_initializer=conv_weight_init, biases_initializer=conv_bias_init,
weights_regularizer=conv_regularizer)
network = slim.conv2d(
network, 32, [3, 3], stride=1, activation_fn=nonlinearity,
padding="SAME", normalizer_fn=batch_norm_fn, scope="conv1_2",
weights_initializer=conv_weight_init, biases_initializer=conv_bias_init,
weights_regularizer=conv_regularizer)
# NOTE(nwojke): This is missing a padding="SAME" to match the CNN
# architecture in Table 1 of the paper. Information on how this affects
# performance on MOT 16 training sequences can be found in
# issue 10 https://github.com/nwojke/deep_sort/issues/10
network = slim.max_pool2d(network, [3, 3], [2, 2], scope="pool1")
network = residual_block(
network, "conv2_1", nonlinearity, conv_weight_init, conv_bias_init,
conv_regularizer, increase_dim=False, is_first=True)
network = residual_block(
network, "conv2_3", nonlinearity, conv_weight_init, conv_bias_init,
conv_regularizer, increase_dim=False)
network = residual_block(
network, "conv3_1", nonlinearity, conv_weight_init, conv_bias_init,
conv_regularizer, increase_dim=True)
network = residual_block(
network, "conv3_3", nonlinearity, conv_weight_init, conv_bias_init,
conv_regularizer, increase_dim=False)
network = residual_block(
network, "conv4_1", nonlinearity, conv_weight_init, conv_bias_init,
conv_regularizer, increase_dim=True)
network = residual_block(
network, "conv4_3", nonlinearity, conv_weight_init, conv_bias_init,
conv_regularizer, increase_dim=False)
feature_dim = network.get_shape().as_list()[-1]
network = slim.flatten(network)
network = slim.dropout(network, keep_prob=0.6)
network = slim.fully_connected(
network, feature_dim, activation_fn=nonlinearity,
normalizer_fn=batch_norm_fn, weights_regularizer=fc_regularizer,
scope="fc1", weights_initializer=fc_weight_init,
biases_initializer=fc_bias_init)
features = network
# Features in rows, normalize axis 1.
features = slim.batch_norm(features, scope="ball", reuse=reuse)
feature_norm = tf.sqrt(
tf.constant(1e-8, tf.float32) +
tf.reduce_sum(tf.square(features), [1], keepdims=True))
features = features / feature_norm
return features, None
def _network_factory(weight_decay=1e-8):
def factory_fn(image, reuse):
with slim.arg_scope([slim.batch_norm, slim.dropout],
is_training=False):
with slim.arg_scope([slim.conv2d, slim.fully_connected,
slim.batch_norm, slim.layer_norm],
reuse=reuse):
features, logits = _create_network(
image, reuse=reuse, weight_decay=weight_decay)
return features, logits
return factory_fn
def _preprocess(image):
image = image[:, :, ::-1] # BGR to RGB
return image
def parse_args():
"""Parse command line arguments.
"""
parser = argparse.ArgumentParser(description="Freeze old model")
parser.add_argument(
"--checkpoint_in",
default="resources/networks/mars-small128.ckpt-68577",
help="Path to checkpoint file")
parser.add_argument(
"--graphdef_out",
default="resources/networks/mars-small128.pb")
return parser.parse_args()
def main():
args = parse_args()
with tf.Session(graph=tf.Graph()) as session:
input_var = tf.placeholder(
tf.uint8, (None, 128, 64, 3), name="images")
image_var = tf.map_fn(
lambda x: _preprocess(x), tf.cast(input_var, tf.float32),
back_prop=False)
factory_fn = _network_factory()
features, _ = factory_fn(image_var, reuse=None)
features = tf.identity(features, name="features")
saver = tf.train.Saver(slim.get_variables_to_restore())
saver.restore(session, args.checkpoint_in)
output_graph_def = tf.graph_util.convert_variables_to_constants(
session, tf.get_default_graph().as_graph_def(),
[features.name.split(":")[0]])
with tf.gfile.GFile(args.graphdef_out, "wb") as file_handle:
file_handle.write(output_graph_def.SerializeToString())
if __name__ == "__main__":
main()
