将字典列表转换为tf数据集

Question

我有一个字典，它已经完全预处理，准备输入到BERT模型中。然而，我正在努力把它变成一个tf.dataset。这是我的数据集的一个元素的样子:print(dataset[0])

{'input_ids': <tf.Tensor: shape=(128,), dtype=int64, numpy= array([  101,   171,   112,  2537, 12293,   131, 11250,   118,   118,
2537, 12293,   131, 11250,  1110,  1126,  1237,  1778,  1326,
1687,  1111,  5957,  1398, 11737,  1118,  8129, 14399,  1105,
3230,  9426, 27277,   119,  1135,  1110,  1103,  1148,  1326,
1872,  4418,  1111,  1115,  1555,   117,  1105,  1103,  1148,
2537, 12293,  1326,  1290,  2537, 12293,   131,  9892,  4803,
1107,  1478,   119,  9617,  4986,   170,  4967,  1196,  1103,
1958,  1104,  1103,  1560,  2537, 12293,  1326,  1105,  2767,
1121,  1103, 21169,  1104,  1103, 18061,  1666,  2672,  2441,
117, 11250, 16001,  1103,  4245,   118,   118,   148,  1979,
1320,  1594,  1229,  1378,  1103,  3039,  1104,  1103,  6684,
11250,   119, 23886,   147,   119, 16218,  1105,  6619, 11679,
19644,  2145,  2867,  1112,  1437, 14627,   102,   171,   112,
1110,  1175,   170,  1207,  2851,   189, 14909,  1326,  1909,
112,   102])>, 'input_mask': <tf.Tensor: shape=(128,), dtype=int64, numpy= array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])>, 'segment_ids': <tf.Tensor: shape=(128,), dtype=int64, numpy= array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])>, 'labels': <tf.Tensor: shape=(), dtype=int64, numpy=1>}

我所需要做的就是把它变成tf.data.Dataset()格式，然而，我似乎无法弄清楚如何使任何可用的from_tensor_slices, from_tensors, from_generator功能与我所拥有的工作。

Answer 1

您可以通过使用pandas(或者您可以模拟to_dict方法的输出)来实现这一点

dataset = tf.data.Dataset.from_tensor_slices(pd.DataFrame.from_dict(records).to_dict(orient="list"))

其中records是字典列表。

Answer 2

我也在为此而挣扎，看到他们的tensorflow_datassets模块在字典中返回数据点(例如https://www.tensorflow.org/tutorials/images/segmentation)并且似乎没有官方功能存在从字典列表创建数据集。

我想出了这个解决方案:

input_ids = tf.data.Dataset.from_tensor_slices([d['input_ids'] for d in dataset])
input_masks = tf.data.Dataset.from_tensor_slices([d['input_mask'] for d in dataset])
segment_ids = tf.data.Dataset.from_tensor_slices([d['segment_ids'] for d in dataset])
labels = tf.data.Dataset.from_tensor_slices([d['labels'] for d in dataset])
ds = tf.data.Dataset.zip((input_ids, input_masks, segment_ids, labels))
ds = ds.map(lambda x, y, z, l: {"input_ids": x, "input_masks": y,
"segment_ids": z, "labels": l}

但是，它不是真正可伸缩的。