如何将两个dict MapDataset合并为一个MapDataset

Question

我在将几个MapDataset连接到一个MapDataset时遇到问题。例如，一个MapDataset是：

<MapDataset element_spec={'input_ids_task1': TensorSpec(), 'mask_task1': TensorSpec(), 'type_ids_task1': TensorSpec()}

另一个是：

<MapDataset element_spec={'input_ids_task2': TensorSpec(), 'mask_task2': TensorSpec(), 'type_ids_task2': TensorSpec()}

我想把它们连接到：

<MapDataset element_spec={'input_ids_task1': TensorSpec(), 'mask_task1': TensorSpec(), 'type_ids_task1': TensorSpec(), 'input_ids_task2': TensorSpec(), 'mask_task2': TensorSpec(), 'type_ids_task2': TensorSpec()}

我已经找到了一些答案，可以通过以下方式压缩这两个数据集：

h = tf.data.Dataset.zip((a, b))

那么h将是ZipDataset:

<ZipDataset element_spec=({'input_ids_task1': TensorSpec(), 'mask_task1': TensorSpec(), 'type_ids_task1': TensorSpec()}, {'input_ids_task2': TensorSpec(), 'mask_task2': TensorSpec(), 'type_ids_task2': TensorSpec()})

因为这两个数据集将是元组中的两个dict。

我可以通过以下方式检索MapDataset：

h.map(lambda x,y: x)

然而，我不确定如何将它们合并为一个dict。

如果这不可能，我可以将输入层更改为包含多个dict的元组以获得数据集输入吗？

Answer 1

不确定你的数据到底是什么样子的，但你应该能够做这样的事情：

import tensorflow as tf
d1 = {
'input_ids_task1': [[1, 2, 3], [1, 2, 2]],
'mask_task1': [[1, 2, 3], [1, 2, 2]],
'type_ids_task1': [[1, 2, 3], [1, 2, 2]] 
}
d2 = {
'input_ids_task2': [[1, 2, 3], [1, 2, 2]],
'mask_task2': [[1, 2, 3], [1, 2, 2]],
'type_ids_task2': [[1, 2, 3], [1, 2, 2]] 
}
dataset1 = tf.data.Dataset.from_tensor_slices((d1))
dataset2 = tf.data.Dataset.from_tensor_slices((d2))
h = tf.data.Dataset.zip((dataset1, dataset2))
h = h.map(lambda x, y: {**x, **y})
print(h)
for d in h:
print(d)

<MapDataset element_spec={'input_ids_task1': TensorSpec(shape=(3,), dtype=tf.int32, name=None), 'mask_task1': TensorSpec(shape=(3,), dtype=tf.int32, name=None), 'type_ids_task1': TensorSpec(shape=(3,), dtype=tf.int32, name=None), 'input_ids_task2': TensorSpec(shape=(3,), dtype=tf.int32, name=None), 'mask_task2': TensorSpec(shape=(3,), dtype=tf.int32, name=None), 'type_ids_task2': TensorSpec(shape=(3,), dtype=tf.int32, name=None)}>
{'input_ids_task1': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, 'mask_task1': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, 'type_ids_task1': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, 'input_ids_task2': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, 'mask_task2': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, 'type_ids_task2': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>}
{'input_ids_task1': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 2], dtype=int32)>, 'mask_task1': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 2], dtype=int32)>, 'type_ids_task1': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 2], dtype=int32)>, 'input_ids_task2': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 2], dtype=int32)>, 'mask_task2': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 2], dtype=int32)>, 'type_ids_task2': <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 2], dtype=int32)>}