def create_lstm_cell(self, batch_size, output_size, state_saver,
state_name):
"""Create the LSTM cell, and initialize state if necessary.
Args:
batch_size: input batch size.
output_size: output size of the lstm cell, [width, height].
state_saver: a state saver object with methods `state` and `save_state`.
state_name: string, the name to use with the state_saver.
Returns:
lstm_cell: the lstm cell unit.
init_state: initial state representations.
step: the step
"""
lstm_cell = lstm_cells.BottleneckConvLSTMCell(
filter_size=(3, 3),
output_size=output_size,
num_units=max(self._min_depth, self._lstm_state_depth),
activation=tf.nn.relu6,
visualize_gates=False)
if state_saver is None:
init_state = lstm_cell.init_state(state_name, batch_size,
tf.float32)
step = None
else:
step = state_saver.state(state_name + '_step')
c = state_saver.state(state_name + '_c')
h = state_saver.state(state_name + '_h')
init_state = (c, h)
return lstm_cell, init_state, step
def test_flatten_state(self):
filter_size = [3, 3]
output_size = [10, 10]
num_units = 15
state_name = 'lstm_state'
batch_size = 4
dtype = tf.float32
unroll = 10
learned_state = False
inputs_large = tf.zeros([4, 10, 10, 5], dtype=tf.float32)
inputs_small = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
cell = lstm_cells.BottleneckConvLSTMCell(filter_size=filter_size,
output_size=output_size,
num_units=num_units,
pre_bottleneck=True,
flatten_state=True)
state = cell.init_state(state_name, batch_size, dtype, learned_state)
for step in range(unroll):
if step % 2 == 0:
inputs = cell.pre_bottleneck(inputs_large, state[1], 0)
else:
inputs = cell.pre_bottleneck(inputs_small, state[1], 1)
output, state = cell(inputs, state)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
output_result, state_result = sess.run([output, state])
self.assertAllEqual((4, 10, 10, 15), output_result.shape)
self.assertAllEqual((4, 10 * 10 * 15), state_result[0].shape)
self.assertAllEqual((4, 10 * 10 * 15), state_result[1].shape)
def test_prebottleneck(self):
filter_size = [3, 3]
output_size = [10, 10]
num_units = 15
state_name = 'lstm_state'
batch_size = 4
dtype = tf.float32
unroll = 10
learned_state = False
inputs_large = tf.zeros([4, 10, 10, 5], dtype=tf.float32)
inputs_small = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
cell = lstm_cells.BottleneckConvLSTMCell(filter_size=filter_size,
output_size=output_size,
num_units=num_units,
pre_bottleneck=True)
state = cell.init_state(state_name, batch_size, dtype, learned_state)
for step in range(unroll):
if step % 2 == 0:
inputs = cell.pre_bottleneck(inputs_large, state[1], 0)
else:
inputs = cell.pre_bottleneck(inputs_small, state[1], 1)
output, state = cell(inputs, state)
self.assertAllEqual([4, 10, 10, 15], output.shape.as_list())
self.assertAllEqual([4, 10, 10, 15], state[0].shape.as_list())
self.assertAllEqual([4, 10, 10, 15], state[1].shape.as_list())
def test_run_lstm_cell(self):
filter_size = [3, 3]
output_size = [10, 10]
num_units = 15
state_name = 'lstm_state'
batch_size = 4
dtype = tf.float32
learned_state = False
inputs = tf.zeros([4, 10, 10, 3], dtype=tf.float32)
cell = lstm_cells.BottleneckConvLSTMCell(filter_size=filter_size,
output_size=output_size,
num_units=num_units)
init_state = cell.init_state(state_name, batch_size, dtype,
learned_state)
output, state_tuple = cell(inputs, init_state)
self.assertAllEqual([4, 10, 10, 15], output.shape.as_list())
self.assertAllEqual([4, 10, 10, 15], state_tuple[0].shape.as_list())
self.assertAllEqual([4, 10, 10, 15], state_tuple[1].shape.as_list())
def test_get_init_learned_state(self):
filter_size = [3, 3]
output_size = [10, 10]
num_units = 15
state_name = 'lstm_state'
batch_size = 4
dtype = tf.float32
learned_state = True
cell = lstm_cells.BottleneckConvLSTMCell(filter_size=filter_size,
output_size=output_size,
num_units=num_units)
init_c, init_h = cell.init_state(state_name, batch_size, dtype,
learned_state)
self.assertEqual(tf.float32, init_c.dtype)
self.assertEqual(tf.float32, init_h.dtype)
self.assertAllEqual([4, 10, 10, 15], init_c.shape.as_list())
self.assertAllEqual([4, 10, 10, 15], init_h.shape.as_list())
def test_get_init_state(self):
filter_size = [3, 3]
output_dim = 10
output_size = [output_dim] * 2
num_units = 15
state_name = 'lstm_state'
batch_size = 4
dtype = tf.float32
learned_state = False
cell = lstm_cells.BottleneckConvLSTMCell(filter_size=filter_size,
output_size=output_size,
num_units=num_units)
init_c, init_h = cell.init_state(state_name, batch_size, dtype,
learned_state)
self.assertEqual(tf.float32, init_c.dtype)
self.assertEqual(tf.float32, init_h.dtype)
with self.test_session() as sess:
init_c_res, init_h_res = sess.run([init_c, init_h])
self.assertAllClose(np.zeros((4, 10, 10, 15)), init_c_res)
self.assertAllClose(np.zeros((4, 10, 10, 15)), init_h_res)
def extract_features(self,
preprocessed_inputs,
state_saver=None,
state_name='lstm_state',
unroll_length=5,
scope=None):
"""Extracts features from preprocessed inputs.
The features include the base network features, lstm features and SSD
features, organized in the following name scope:
<parent scope>/MobilenetV1/...
<parent scope>/LSTM/...
<parent scope>/FeatureMaps/...
Args:
preprocessed_inputs: A [batch, height, width, channels] float tensor
representing a batch of consecutive frames from video clips.
state_saver: A state saver object with methods `state` and `save_state`.
state_name: A python string for the name to use with the state_saver.
unroll_length: The number of steps to unroll the lstm.
scope: The scope for the base network of the feature extractor.
Returns:
A list of tensors where the ith tensor has shape [batch, height_i,
width_i, depth_i]
"""
preprocessed_inputs = shape_utils.check_min_image_dim(
33, preprocessed_inputs)
with slim.arg_scope(
mobilenet_v1.mobilenet_v1_arg_scope(
is_training=self._is_training)):
with (slim.arg_scope(self._conv_hyperparams_fn())
if self._override_base_feature_extractor_hyperparams else
context_manager.IdentityContextManager()):
with slim.arg_scope([slim.batch_norm], fused=False):
# Base network.
with tf.variable_scope(scope,
self._base_network_scope,
reuse=self._reuse_weights) as scope:
net, image_features = mobilenet_v1.mobilenet_v1_base(
ops.pad_to_multiple(preprocessed_inputs,
self._pad_to_multiple),
final_endpoint='Conv2d_13_pointwise',
min_depth=self._min_depth,
depth_multiplier=self._depth_multiplier,
scope=scope)
with slim.arg_scope(self._conv_hyperparams_fn()):
with slim.arg_scope([slim.batch_norm],
fused=False,
is_training=self._is_training):
# ConvLSTM layers.
with tf.variable_scope(
'LSTM', reuse=self._reuse_weights) as lstm_scope:
lstm_cell = lstm_cells.BottleneckConvLSTMCell(
filter_size=(3, 3),
output_size=(net.shape[1].value, net.shape[2].value),
num_units=max(self._min_depth, self._lstm_state_depth),
activation=tf.nn.relu6,
visualize_gates=True)
net_seq = list(tf.split(net, unroll_length))
if state_saver is None:
init_state = lstm_cell.init_state(
state_name, net.shape[0].value / unroll_length,
tf.float32)
else:
c = state_saver.state('%s_c' % state_name)
h = state_saver.state('%s_h' % state_name)
init_state = (c, h)
# Identities added for inputing state tensors externally.
c_ident = tf.identity(init_state[0],
name='lstm_state_in_c')
h_ident = tf.identity(init_state[1],
name='lstm_state_in_h')
init_state = (c_ident, h_ident)
net_seq, states_out = rnn_decoder.rnn_decoder(
net_seq, init_state, lstm_cell, scope=lstm_scope)
batcher_ops = None
self._states_out = states_out
if state_saver is not None:
self._step = state_saver.state('%s_step' % state_name)
batcher_ops = [
state_saver.save_state('%s_c' % state_name,
states_out[-1][0]),
state_saver.save_state('%s_h' % state_name,
states_out[-1][1]),
state_saver.save_state('%s_step' % state_name,
self._step - 1)
]
with tf_ops.control_dependencies(batcher_ops):
image_features['Conv2d_13_pointwise_lstm'] = tf.concat(
net_seq, 0)
# Identities added for reading output states, to be reused externally.
tf.identity(states_out[-1][0], name='lstm_state_out_c')
tf.identity(states_out[-1][1], name='lstm_state_out_h')
# SSD layers.
with tf.variable_scope('FeatureMaps',
reuse=self._reuse_weights):
feature_maps = feature_map_generators.multi_resolution_feature_maps(
feature_map_layout=self._feature_map_layout,
depth_multiplier=(self._depth_multiplier),
min_depth=self._min_depth,
insert_1x1_conv=True,
image_features=image_features)
return feature_maps.values()