将 tensorflow 2.0 BatchDataset 转换为 numpy 数组

Question

我有这个代码：

(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(TRAIN_BUF).batch(BATCH_SIZE)
test_dataset = tf.data.Dataset.from_tensor_slices(test_images).shuffle(TRAIN_BUF).batch(BATCH_SIZE)

print(train_dataset, type(train_dataset), test_dataset, type(test_dataset))

我想将这两个BatchDataset 变量转换为numpy arrays，我可以轻松做到吗？我正在使用TF 2.0，但我刚刚找到了将tf.data 与TF 1.0 一起转换的代码

Answer 1

对数据集进行批处理后，最后一批的形状可能与其余批次的形状不同。例如，如果您的数据集中总共有 100 个元素，并且您的批次大小为 6，则最后一批的大小仅为 4。（100 = 6 * 16 + 4）。

因此，在这种情况下，您将无法直接将数据集转换为 numpy。因此，您必须在批处理方法中使用 drop_remainder 参数为 True。如果大小不正确，它将丢弃最后一批。

之后，我附上了如何将数据集转换为 Numpy 的代码。

import tensorflow as tf
import numpy as np

(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()

TRAIN_BUF=1000
BATCH_SIZE=64

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).
                          shuffle(TRAIN_BUF).batch(BATCH_SIZE, drop_remainder=True)
test_dataset = tf.data.Dataset.from_tensor_slices(test_images).
                          shuffle(TRAIN_BUF).batch(BATCH_SIZE, drop_remainder=True)

# print(train_dataset, type(train_dataset), test_dataset, type(test_dataset))

train_np = np.stack(list(train_dataset))
test_np = np.stack(list(test_dataset))
print(type(train_np), train_np.shape)
print(type(test_np), test_np.shape)

输出：

<class 'numpy.ndarray'> (937, 64, 28, 28)
<class 'numpy.ndarray'> (156, 64, 28, 28)