def custom_pack_batch(x):
"""Map-function."""
(k1_packed, k1_segmentation, k1_position, k2_packed, k2_segmentation,
k2_position) = (
pack_sequences_ops.pack_sequences2(
# cast to int64 for compatibility with custom ops
tf.cast(x[k1], tf.int64),
tf.cast(x[k2], tf.int64),
length[k1],
length[k2]))
packed = {
k1: k1_packed,
k1 + "_segmentation": k1_segmentation,
k1 + "_position": k1_position,
}
if len(keys) == 2:
packed.update({
k2: k2_packed,
k2 + "_segmentation": k2_segmentation,
k2 + "_position": k2_position,
})
# cast back to int32
for k, v in packed.items():
packed[k] = tf.cast(v, tf.int32)
return packed
def map_fn_custom(x):
"""Map-function."""
(k1_packed, k1_segmengation, k1_position,
k2_packed, k2_segmentation, k2_position) = (
pack_sequences_ops.pack_sequences2(x[k1], x[k2], length))
packed = {
k1: k1_packed,
k1 + "_segmentation": k1_segmengation,
k1 + "_position": k1_position,
k2: k2_packed,
k2 + "_segmentation": k2_segmentation,
k2 + "_position": k2_position,
}
return tf.data.Dataset.from_tensor_slices(packed)
def map_fn_custom(x):
"""Map-function."""
(k1_packed, k1_segmengation, k1_position, k2_packed, k2_segmentation,
k2_position) = (pack_sequences_ops.pack_sequences2(
x[k1], x[k2], length[k1], length[k2]))
packed = {
k1: k1_packed,
k1 + "_segmentation": k1_segmengation,
k1 + "_position": k1_position,
}
if len(keys) == 2:
packed.update({
k2: k2_packed,
k2 + "_segmentation": k2_segmentation,
k2 + "_position": k2_position,
})
return packed
def test_random_inputs(self):
for _ in range(10):
batch_size = np.random.randint(900, 1100, size=[])
input_seqlen = np.random.randint(1, 10, size=[])
target_seqlen = np.random.randint(1, 10, size=[])
inputs_list = []
targets_list = []
for _ in range(batch_size):
input_num_pads = np.random.randint(0, input_seqlen, size=[])
input_pads = np.full([input_num_pads], 0, dtype=np.int32)
inputs = np.random.randint(
1, 10, size=[input_seqlen - input_num_pads])
inputs = np.concatenate([inputs, input_pads], axis=0)
target_num_pads = np.random.randint(0, target_seqlen, size=[])
target_pads = np.full([target_num_pads], 0, dtype=np.int32)
targets = np.random.randint(
1, 10, size=[target_seqlen - target_num_pads])
targets = np.concatenate([targets, target_pads], axis=0)
inputs_list.append(inputs)
targets_list.append(targets)
input_maxlen = np.random.randint(input_seqlen,
input_seqlen + 10,
size=[])
target_maxlen = np.random.randint(target_seqlen,
target_seqlen + 10,
size=[])
(inputs_packed2, inputs_segmentation2, inputs_positions2,
targets_packed2, targets_segmentation2,
targets_positions2) = (pack_sequences_ops.pack_sequences2(
inputs_list, targets_list, input_maxlen, target_maxlen))
(inputs_packed_k, inputs_segmentation_k, inputs_positions_k,
targets_packed_k, targets_segmentation_k,
targets_positions_k) = (_pack_sequences_k(inputs_list,
targets_list,
input_maxlen,
target_maxlen))
self.assertAllEqual(inputs_packed2, inputs_packed_k)
self.assertAllEqual(inputs_segmentation2, inputs_segmentation_k)
self.assertAllEqual(inputs_positions2, inputs_positions_k)
self.assertAllEqual(targets_packed2, targets_packed_k)
self.assertAllEqual(targets_segmentation2, targets_segmentation_k)
self.assertAllEqual(targets_positions2, targets_positions_k)
def test_pack_sequences2(self):
inputs = [
[1, 2, 3],
[4, 5, 0],
[6, 0, 0],
]
targets = [
[10, 0, 0],
[20, 30, 40],
[50, 60, 0],
]
max_length = 5
(inputs_packed, inputs_segmentation, inputs_position, targets_packed,
targets_segmentation,
targets_position) = (pack_sequences_ops.pack_sequences2(
inputs, targets, max_length, max_length))
self.assertAllEqual(inputs_packed, [
[1, 2, 3, 4, 5],
[6, 0, 0, 0, 0],
])
self.assertAllEqual(inputs_segmentation, [
[1, 1, 1, 2, 2],
[1, 0, 0, 0, 0],
])
self.assertAllEqual(inputs_position, [
[0, 1, 2, 0, 1],
[0, 0, 0, 0, 0],
])
self.assertAllEqual(targets_packed, [
[10, 20, 30, 40, 0],
[50, 60, 0, 0, 0],
])
self.assertAllEqual(targets_segmentation, [
[1, 2, 2, 2, 0],
[1, 1, 0, 0, 0],
])
self.assertAllEqual(targets_position, [
[0, 0, 1, 2, 0],
[0, 1, 0, 0, 0],
])