【Tensorflow 2.0 正式版教程】tf.data.Dataset的基本使用方法

Tensorflow 2.0中提供了专门用于数据输入的接口tf.data.Dataset，可以简洁高效的实现数据的读入、打乱（shuffle）、增强（augment）等功能。下面以一个简单的实例讲解该功能的基本使用方法。

首先手工创建一个非常简单的数据集，该数据包含10个样本，每个样本由1个浮点数组成。

data = np.array([0.1, 0.4, 0.6, 0.2, 0.8, 0.8, 0.4, 0.9, 0.3, 0.2])

其中大于0.5的样本为正样本，即标签记为1，否则为0。

label = np.array([0, 0, 1, 0, 1, 1, 0, 1, 0, 0])

然后，可以通过tf.data.Dataset.from_tensor_slices建立数据集。

dataset = tf.data.Dataset.from_tensor_slices((data, label))

该数据集可以直接由python原生语法进行迭代

for x, y in dataset:
	print(x, y)

可以得到如下输出

tf.Tensor(0.1, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.4, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.6, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.2, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.8, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.8, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.4, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.9, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.3, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.2, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)

但是，更多情况下训练网络时数据不止迭代一轮，可以通过执行repeat()使数据集能多次迭代，这种写法下推荐生成迭代器自行迭代。

dataset = dataset.repeat()
it = dataset.__iter__()
for i in range(20):
    x, y = it.next()
    print(x, y)

shuffle()是随机打乱样本次序，参数buffer_size建议设为样本数量，过大会浪费内存空间，过小会导致打乱不充分。

dataset = dataset.shuffle(buffer_size=10)
it = dataset.__iter__()
for i in range(10):
    x, y = it.next()
    print(x, y)

batch()是使迭代器一次获取多个样本

dataset_batch = dataset.batch(batch_size=5)
it = dataset_batch.__iter__()
for i in range(2):
    x, y = it.next()
    print(x, y)

此时的输出格式为

tf.Tensor([0.9 0.1 0.3 0.2 0.8], shape=(5,), dtype=float64) tf.Tensor([1 0 0 0 1], shape=(5,), dtype=int64)
tf.Tensor([0.4 0.4 0.8 0.8 0.4], shape=(5,), dtype=float64) tf.Tensor([0 0 1 1 0], shape=(5,), dtype=int64)

最后，介绍map()这一核心函数。该函数的输入参数map_func应为一个函数，在该函数中实现我们需要的对数据的变换。具体应用场景如图片加载、数据增强、标签one hot化等。下面以one hot化和添加噪声为例具体说明。

one hot化的函数实现如下

def one_hot(x, y):
    if y == 0:
        return x, np.array([1, 0])
    else:
        return x, np.array([0, 1])

数据集对应执行

dataset_one_hot = dataset.map(one_hot)
it = dataset_one_hot.__iter__()
for i in range(10):
    x, y = it.next()
    print(x, y)

此时的输出为

tf.Tensor(0.8, shape=(), dtype=float64) tf.Tensor([0 1], shape=(2,), dtype=int64)
tf.Tensor(0.6, shape=(), dtype=float64) tf.Tensor([0 1], shape=(2,), dtype=int64)
tf.Tensor(0.4, shape=(), dtype=float64) tf.Tensor([1 0], shape=(2,), dtype=int64)
tf.Tensor(0.1, shape=(), dtype=float64) tf.Tensor([1 0], shape=(2,), dtype=int64)
tf.Tensor(0.4, shape=(), dtype=float64) tf.Tensor([1 0], shape=(2,), dtype=int64)
tf.Tensor(0.1, shape=(), dtype=float64) tf.Tensor([1 0], shape=(2,), dtype=int64)
tf.Tensor(0.9, shape=(), dtype=float64) tf.Tensor([0 1], shape=(2,), dtype=int64)
tf.Tensor(0.8, shape=(), dtype=float64) tf.Tensor([0 1], shape=(2,), dtype=int64)
tf.Tensor(0.2, shape=(), dtype=float64) tf.Tensor([1 0], shape=(2,), dtype=int64)
tf.Tensor(0.2, shape=(), dtype=float64) tf.Tensor([1 0], shape=(2,), dtype=int64)

对数据进行固定形式上的变化，可将函数直接作为参数输入。但是，包含随机信息的数据变化则需要tf.py_function辅助实现，如数据增强中数据添加随机噪声、图像的随机翻转都属于包含随机信息。

def add_noise(x, y):
    x += np.random.uniform(0.0, 0.01)
return x, y

如果直接输入

dataset_add_noise = dataset.map(add_noise)

此时迭代结果为

tf.Tensor(0.9040774445322317, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.6040774445322317, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.2040774445322317, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.6040774445322317, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.4040774445322317, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.4040774445322317, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.4040774445322317, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.3040774445322317, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.8040774445322317, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.8040774445322317, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)

可以看到对于每个样本添加的‘随机’噪声是相同的，正确的实现方法应为

dataset_add_noise = dataset.map(lambda x, y: tf.py_function(add_noise, inp=[x, y], Tout=[tf.float64, tf.int64]))

从而可以得到期望的随机结果

tf.Tensor(0.1016960895379668, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.8054602820484098, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.609566791084528, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.40492143124755076, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.2070627361984957, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.8044649643628696, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.2099414947814727, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.8081644487759416, shape=(), dtype=float64) tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(0.4072158822338304, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(0.2041248890575763, shape=(), dtype=float64) tf.Tensor(0, shape=(), dtype=int64)

在map()函数中，还有个很重要的参数num_parallel_calls，可以将数据加载与变换过程并行到多个CPU线程上。由于python语言本身的全局解释锁，想要实现真正的并行计算是非常困难的，所以这个参数实际上非常实用，通常的使用情景是网络训练时，GPU做模型运算的同时CPU加载数据。
还可以直接设置num_parallel_calls=tf.data.experimental.AUTOTUNE，这样会自动设置为最大的可用线程数，机器算力拉满。

完整的代码可以在这里找到
https://github.com/Apm5/tensorflow_2.0_tutorial/blob/master/data/dataset.py