def test_shuffle_batches(self):
# Shape = [3, 2].
tensor_1 = tf.constant([[1, 2], [3, 4], [5, 6]])
tensor_2 = tf.constant([[11, 12], [13, 14], [15, 16]])
tensor_3 = tf.constant([[21, 22], [23, 24], [25, 26]])
shuffled_tensor_1, shuffled_tensor_2, shuffled_tensor_3 = (
data_utils.shuffle_batches([tensor_1, tensor_2, tensor_3]))
tensor_diff_21 = shuffled_tensor_2 - shuffled_tensor_1
tensor_diff_31 = shuffled_tensor_3 - shuffled_tensor_1
self.assertAllEqual(tensor_diff_21, [[10, 10], [10, 10], [10, 10]])
self.assertAllEqual(tensor_diff_31, [[20, 20], [20, 20], [20, 20]])
def train_one_iteration(inputs):
"""Trains the model for one iteration.
Args:
inputs: A dictionary for training inputs.
Returns:
The training loss for this iteration.
"""
_, side_outputs = pipelines.create_model_input(
inputs, FLAGS.model_input_keypoint_type, keypoint_profile_2d,
keypoint_profile_3d)
keypoints_2d = side_outputs[common_module.KEY_PREPROCESSED_KEYPOINTS_2D]
keypoints_3d, _ = keypoint_preprocessor_3d(
inputs[common_module.KEY_KEYPOINTS_3D],
keypoint_profile_3d,
normalize_keypoints_3d=True)
keypoints_2d, keypoints_3d = data_utils.shuffle_batches(
[keypoints_2d, keypoints_3d])
return model.train((keypoints_2d, keypoints_3d), **optimizers)