使用tf.data.Dataset的Tensorflow性能下降

Question

我已经在tensorflow中实现了一个简单的训练器类。我正在进行一些实验来检查代码性能，但是在了解 tf.data.Dataset 和 tf.function 的作用下遇到了问题。

在下面，我将介绍我已经运行的测试，最后，关于我得到的结果会有一些疑问。

配置：Intel i3 cpu，tensorflow-cpu 2.1

class Trainer:
    def __init__(self, model, optimizer, loss):
        self.model = model
        self.loss_function = loss
        self.optimizer = optimizer

    @tf.function
    def train_step(self, inputs, targets):
        with tf.GradientTape() as tape:
            predictions = self.model(inputs)
            loss = self.loss_function(targets, predictions)
        gradients = tape.gradient(loss, self.model.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
        return loss

    # fit using dataset
    @tf.function
    def fit0(self, dataset, epochs):
        for epoch in tf.range(epochs):
            for input_batch, target_batch in dataset:
                self.train_step(input_batch, target_batch)

    # fit using list of tensors
    @tf.function
    def fit1(self, inputs, targets, epochs):
        for epoch in tf.range(epochs):
            for input_batch, target_batch in zip(inputs, targets):
                self.train_step(input_batch, target_batch)

在以下 train_step 中，将始终包装在 tf.function 中。

fit0 ， fit1 将在有和没有 tf。功能的情况下进行测试。

这是运行测试的代码：

input_size = 10000
batch_size = 100
q = input_size // batch_size

# create random inputs (x) and outputs (y)
x = tf.random.normal((input_size, 1), dtype=tf.float32)
y = tf.random.normal((input_size, 1), dtype=tf.float32)

splits = tf.fill([q, ], batch_size)

# create a list of tensors rappresenting batches
x_list = tf.split(x, splits)
y_list = tf.split(y, splits)

# create datasets in the different ways
dataset0 = tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size)
dataset1 = tf.data.Dataset.from_tensor_slices((tf.stack(x_list), tf.stack(y_list)))

# model definition
model = tf.keras.Sequential([
    tf.keras.layers.Dense(20, activation='tanh', input_shape=(1,)),
    tf.keras.layers.Dense(1, activation='linear')])

# trainer initialization
trainer = Trainer(model=model, optimizer=tf.keras.optimizers.Adam(), loss=tf.keras.losses.MeanSquaredError())

# first run to perform initializations
time0 = time.time()
trainer.fit0(dataset=dataset0, epochs=tf.constant(1, dtype=tf.int32))
time0 = time.time() - time0

time1 = time.time()
trainer.fit0(dataset=dataset1, epochs=tf.constant(1, dtype=tf.int32))
time1 = time.time() - time1

time2 = time.time()
trainer.fit1(inputs=x_list, targets=y_list, epochs=tf.constant(1, dtype=tf.int32))
time2 = time.time() - time2

print("first fit0 with dataset0 took {} seconds".format(time0))
print("first fit0 with dataset1 took {} seconds".format(time1))
print("first fit1 with tensorlist took {} seconds".format(time2))

# measure performances
time0 = time.time()
trainer.fit0(dataset=dataset0, epochs=tf.constant(100, dtype=tf.int32))
time0 = time.time() - time0

time1 = time.time()
trainer.fit0(dataset=dataset1, epochs=tf.constant(100, dtype=tf.int32))
time1 = time.time() - time1

time2 = time.time()
trainer.fit1(inputs=x_list, targets=y_list, epochs=tf.constant(100, dtype=tf.int32))
time2 = time.time() - time2

print("fit0 with dataset0 took {} seconds".format(time0))
print("fit0 with dataset1 took {} seconds".format(time1))
print("fit1 with tensorlist took {} seconds".format(time2))

这里是测试结果：

第一个测试是100批，每批100个样品。

input_size = 10000
batch_size = 100

没有@ tf.function：
数据集0的第一个fit0花了0.9953532218933105秒
数据集1的第一个fit0花费了0.07995295524597168秒
张量表的首次拟合1花费了0.05196571350097656秒
带有数据集0的fit0用了10.46957802772522秒
具有数据集1的fit0花费了7.822799205780029秒
带有张量表的fit1花费了4.650130748748779秒

具有@ tf.function：
数据集0的第一个fit0用了1.4042332172393799秒
数据集1的第一个fit0用了0.46071624755859375秒
张量表的第一次fit1用了7.3524699211120605秒
fit0与dataset0花费了15.077088832855225秒
具有数据集1的fit0花费了9.136569738388062秒
张量表的fit1花费了2.1366817951202393秒

第二个是一批100000个样本。

input_size = 100000
batch_size = 100000

没有@ tf.function：
数据集0的第一个fit0用了1.1792669296264648秒
数据集1的第一个fit0花费了0.027983427047729492秒
张量表的首次拟合1花费了0.020987749099731445秒
fit0与dataset0花费了28.71895956993103秒
具有数据集1的fit0用了2.730872869491577秒
tensorlist的fit1用了2.194814682006836秒

具有@ tf.function：
数据集0的第一个fit0用了1.5979444980621338秒
数据集1的第一个fit0用了0.4557182788848877秒
张量表的第一次fit1用了0.3708038330078125秒
具有数据集0的fit0用了36.43854784965515秒
具有数据集1的fit0花费了9.819332122802734秒
具有张量表的fit1花费了2.1136972904205322秒

问题：

1. 为什么 tf.data.Dataset 用 tf.function 包装后性能最差？
2. 即使dataset0和dataset1在功能上等效。两者之间的内在区别是什么？为什么dataset1的性能优于dataset0？
具有 tf.function 的
3. fit1 具有最佳的长期性能。
   • 使用 tf.data.Dataset 是否可以达到相同的性能？
   • 为什么要花这么多时间进行初始化？
     如果使用100个批次，则第一次运行耗时7.3524699211120605秒，而这一次通过增加批次数而增加。
     我猜是因为亲笔签名正在创建更大的图，展开了不同批次的计算。不过，我看不到任何并行化的机会，因为每一批都依赖于上一批的结果。

Answer 1

我的性能得到了很好的改善，代码和结果如下所示。
但是，我只能部分回答这些问题，尤其是第二个问题仍然开放。

配置：Intel i3 cpu，tensorflow-cpu 2.1
这是功能 fit0 的改进代码，其他 Trainer 类未更改：

# fit using dataset
@tf.function
def fit0(self, dataset, epochs, batches, unroll=1):
    tf.assert_equal(tf.is_tensor(unroll), False, "unroll must be a python variable.")
    tf.assert_equal(tf.math.floormod(batches, unroll), tf.constant(0), "unroll must be a divisor of batches.")

    entries = epochs * batches / unroll
    it = iter(dataset)

    for entry in tf.range(entries):
        # this loop gets unrolled if unroll
        # is python variable, not a tensor.
        for _ in range(unroll):
            input_batch, target_batch = next(it)
            self.train_step(input_batch, target_batch)

这是运行测试的代码：

input_size = 100000
batch_size = 100
num_epochs = 100
num_unroll = 5

num_batches = input_size // batch_size

# create random inputs (x) and outputs (y)
x = tf.random.normal((input_size, 1), dtype=tf.float32)
y = tf.random.normal((input_size, 1), dtype=tf.float32)

splits = tf.fill([num_batches, ], batch_size)
x_list, y_list = tf.split(x, splits), tf.split(y, splits)

# create dataset
dataset0 = tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size).cache().prefetch(1).repeat(num_epochs)
dataset1 = tf.data.Dataset.from_tensor_slices((tf.stack(x_list), tf.stack(y_list))).cache().prefetch(1).repeat(num_epochs)

# model definition
model = tf.keras.Sequential([
    tf.keras.layers.Dense(20, activation='tanh', input_shape=(1,)),
    tf.keras.layers.Dense(1, activation='linear')])

# trainer initialization
trainer = Trainer(model=model, optimizer=tf.keras.optimizers.Adam(), loss=tf.keras.losses.MeanSquaredError())

# first run to perform initializations
time0 = time.perf_counter()
trainer.fit0(
    dataset=dataset0,
    epochs=tf.constant(1, dtype=tf.int32),
    batches=tf.constant(num_batches, dtype=tf.int32),
    unroll=num_unroll)
time0 = time.perf_counter() - time0

time1 = time.perf_counter()
trainer.fit0(
    dataset=dataset1,
    epochs=tf.constant(1, dtype=tf.int32),
    batches=tf.constant(num_batches, dtype=tf.int32),
    unroll=num_unroll)
time1 = time.perf_counter() - time1

time2 = time.perf_counter()
trainer.fit1(inputs=x_list, targets=y_list, epochs=tf.constant(1, dtype=tf.int32))
time2 = time.perf_counter() - time2

print("first fit0 with dataset0 took {} seconds".format(time0))
print("first fit0 with dataset1 took {} seconds".format(time1))
print("first fit1 with tensorlist took {} seconds".format(time2))

# measure performances
time0 = time.perf_counter()
trainer.fit0(
    dataset=dataset0,
    epochs=tf.constant(num_epochs, dtype=tf.int32),
    batches=tf.constant(num_batches, dtype=tf.int32),
    unroll=num_unroll)
time0 = time.perf_counter() - time0

time1 = time.perf_counter()
trainer.fit0(
    dataset=dataset1,
    epochs=tf.constant(num_epochs, dtype=tf.int32),
    batches=tf.constant(num_batches, dtype=tf.int32),
    unroll=num_unroll)
time1 = time.perf_counter() - time1

time2 = time.perf_counter()
trainer.fit1(inputs=x_list, targets=y_list, epochs=tf.constant(num_epochs, dtype=tf.int32))
time2 = time.perf_counter() - time2

print("fit0 with dataset0 took {} seconds".format(time0))
print("fit0 with dataset1 took {} seconds".format(time1))
print("fit1 with tensorlist took {} seconds".format(time2))

1. 为什么 tf.data.Dataset 用 tf.function 包装后性能最差？

我不知道幕后到底发生了什么，但是可以通过替换以下内容来解决：

dataset0 = tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size)
dataset1 = tf.data.Dataset.from_tensor_slices((tf.stack(x_list), tf.stack(y_list)))

有了这个新的数据集，其中还包括历程以及使用缓存和预取。

dataset0 = tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size).cache().prefetch(1).repeat(num_epochs)
dataset1 = tf.data.Dataset.from_tensor_slices((tf.stack(x_list), tf.stack(y_list))).cache().prefetch(1).repeat(num_epochs)

可以找到其他更多信息here。

我测试了 fit0 ， fit1 ，有和没有 tf.function ，但通过进行此更改，我始终可以通过使用获得更好的性能tf.function ，因此仅显示后者。

使用的input_size大10倍。这是测试结果：

第一个测试是对1000个批次进行测试，每个批次100个样品。
请注意，与num_unroll = 1相比，num_unroll = 5可以提高性能。设置num_unroll> 5不会提供任何进一步的改进。

input_size = 100000
batch_size = 100
num_epochs = 100
num_unroll = 5

数据集0的首次拟合0花费了2.2224882999999993秒
数据集1的第一个fit0用了0.804360700000001秒
张量表的第一次fit1花费了88.2123332秒
使用数据集0的fit0用了35.27911590000001秒
具有数据集1的fit0用了20.370243099999982秒
具有张量列表的fit1花费了23.66727979999999秒

第二个是带有1批1000000个样本的

input_size = 100000
batch_size = 100000

input_size = 1000000
batch_size = 1000000
num_epochs = 100
num_unroll = 1

数据集为0的第一个fit0用了4.3616363秒
数据集1的第一个fit0用了0.7977632000000003秒
张量表的第一次fit1用了0.7329889000000005秒
fit0与dataset0花费了21.131495899999997秒
数据集1的fit0用了19.915148600000002秒
tensorlist的fit1花费了19.817472700000003秒

以上结果可以回答这个问题：

具有 tf.function 的
1. fit1 具有最佳的长期性能。
   • 使用 tf.data.Dataset 是否可以达到相同的性能？
   • 为什么要花这么多时间进行初始化？
     当使用100批次时，第一次运行耗时7.3524699211120605秒，因此 通过增加批数来增加时间。我想是 因为签名会创建更大的图，所以展开 计算不同的批次。我看不到任何机会 但是并行化，因为每个批次都取决于结果 上一个。

通过检查 TensorBoard 上图形的结构，可以很容易地看到，在 fit1 函数上使用签名，可以通过完全展开循环来创建非常大的图形。 这样可以提供更好的性能，但是创建图形的时间很长，而且很可能会过度使用内存，这使得它无法用于更复杂的问题。
但是，如上所示，使用 tf.data.Dataset 可以实现相同的性能，而只需要展开几个循环，从而可以改善图形大小。