Набор данных Tensorflow: получить массив пустых меток

Question

У меня есть tf.data.Dataset, dataset, который имеет описание этой функции:

    feature_description = {
        'text_features': tf.FixedLenFeature([100], tf.int64),
        'numeric_features': tf.FixedLenFeature([200], tf.float32),
        'label': tf.FixedLenFeature([1], tf.int64),
    }

Я хотел бы получить массив NumPy только метки из каждого образца.Я могу получить полный массив NumPy, выполнив следующее:

def load_dataset(dataset):
    """ Load an entire tf dataset into memory
    """
    max_elems = np.iinfo(np.int32).max

    # Make a single batch out of the entire dataset and get that element
    dataset = dataset.batch(max_elems)
    dataset_tensors = tf.contrib.data.get_single_element(dataset)

    # Create a session and evaluate `whole_dataset_tensors` to get arrays.
    with tf.Session() as sess:
        return sess.run(dataset_tensors)

Но это загружает в память полный dataset как массив NumPy (и вызывает на моем ноутбуке OutOfMemoryError).Я просто хочу получить ярлыки.

Одна мысль: может быть, я могу сделать что-то вроде:

dataset = dataset.map(lambda x: x['label'] result = load_dataset(dataset)?

Есть предложения?