def test_sequence_parsing(self, batch_size):
file_pattern = os.path.join(FLAGS.test_tmpdir, 'test.tfrecord')
sequence_length = 3
if not os.path.exists(file_pattern):
self._write_test_sequence_examples(sequence_length, file_pattern)
dataset = tfdata.parallel_read(file_patterns=file_pattern)
# Features
state_spec_1 = tensorspec_utils.ExtendedTensorSpec(
shape=(TEST_IMAGE_SHAPE),
dtype=tf.uint8,
is_sequence=True,
name='image_sequence_feature',
data_format='JPEG')
state_spec_2 = tensorspec_utils.ExtendedTensorSpec(
shape=(2),
dtype=tf.float32,
is_sequence=True,
name='sequence_feature')
feature_tspec = PoseEnvFeature(state=state_spec_1, action=state_spec_2)
feature_tspec = tensorspec_utils.add_sequence_length_specs(
feature_tspec)
# Labels
reward_spec = tensorspec_utils.ExtendedTensorSpec(
shape=(),
dtype=tf.int64,
is_sequence=False,
name='context_feature')
label_tspec = PoseEnvLabel(reward=reward_spec)
label_tspec = tensorspec_utils.add_sequence_length_specs(label_tspec)
dataset = dataset.batch(batch_size, drop_remainder=True)
dataset = tfdata.serialized_to_parsed(dataset, feature_tspec,
label_tspec)
features, labels = dataset.make_one_shot_iterator().get_next()
# Check tensor shapes.
self.assertAllEqual([batch_size, None] + TEST_IMAGE_SHAPE,
features.state.shape.as_list())
self.assertAllEqual([batch_size, None, 2],
features.action.shape.as_list())
self.assertAllEqual([batch_size],
features.state_length.shape.as_list())
self.assertAllEqual([batch_size],
features.action_length.shape.as_list())
self.assertAllEqual([batch_size], labels.reward.shape.as_list())
with self.session() as session:
features_, labels_ = session.run([features, labels])
# Check that images are equal.
for i in range(3):
img = TEST_IMAGE * i
self.assertAllEqual(img, features_.state[0, i])
# Check that numpy shapes are equal.
self.assertAllEqual([batch_size, sequence_length] +
TEST_IMAGE_SHAPE, features_.state.shape)
self.assertAllEqual([sequence_length] * batch_size,
features_.state_length)
self.assertAllEqual([batch_size, sequence_length, 2],
features_.action.shape)
self.assertAllEqual([batch_size], labels_.reward.shape)
def test_add_sequence_length_specs(self):
input_spec = utils.TensorSpecStruct(image1=D1, actions=S7)
modified_spec = utils.add_sequence_length_specs(input_spec)
expected_length_spec = utils.ExtendedTensorSpec(
shape=(), dtype=tf.int64, name='sequence_actions_length')
self.assertEqual(modified_spec.actions_length, expected_length_spec)
