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], ])