tensorflow - 如何通过特定键对张量流中的数据集进行分组，同时避免内存不足？

Question

我使用 tf.dataset 从磁盘（HDFS）加载数据，然后尝试按特定键对数据进行分组。原始数据是这样的：</p>

[["x","a",1],
["y","b",2],
["x","b",2],
["x","b",3],
["x","a",3],
["y","b",4]]

转换后的数据预计为：

[[["x","a",1]
["x","a",3]],
[["x","b",2]
["x","b",3]],
[["y","b",2],
["y","b",4]]]

编码：

dataset = dataset.apply(
                tf.data.experimental.group_by_window(
                    key_func = lambda x : tf.squeeze(
                        tf.string_to_hash_bucket_strong(
                            tf.strings.join([x['key_1'], x['key_2']]),
                            1500000000,
                            [1,2]
                        ),
                        axis=-1
                    ),
                    reduce_func = lambda key, x: x.padded_batch(batch_size=self._window_size, padded_shapes=self._dataset_shapes(dataset), padding_values=self._dataset_default_values(dataset)),
                    window_size =self._window_size
                )
            )

它会产生OOM错误。问题是由 tf.dataset.apply() 或 tf.data.experimental.group_by_window() 引起的，因为当我删除此处的代码时一切正常。我有两个问题：为什么会导致 OOM 以及如何通过特定键对数据进行分组以避免 OOM？

标签： tensorflowout-of-memorytensorflow-datasets