从非常大的 TensorFlow 数据集中仅提取标签

Question

我有一个 120GB 的大型数据集，像这样使用 TensorFlow 2.4.0 加载：

ds_train, ds_test = tfds.load('myds', split=['train[20%:]','train[:20%]'])

由于它的大小，我无法缓存它。

数据集不平衡，我想用sklearn.utils.class_weight.compute_class_weight()计算 class 权重矩阵。 该方法需要标签作为我尝试生成的np.array ，如下所示：

class_labels = np.array([x[1].numpy() for x in list(ds_train)])

这种方法的问题是，它试图将整个数据集加载到失败的 memory 中。 有没有办法从大型tf.data.Dataset中提取标签？