Tensorflow 使用 tf.data.Dataset 降低性能

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.function進行測試。

這是我運行測試的代碼：

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 個樣品。

輸入大小 = 10000
批量大小 = 100

沒有@tf.function：
第一次 fit0 與 dataset0 花費了 0.9953532218933105 秒
第一次 fit0 與 dataset1 花費了 0.07995295524597168 秒
第一個帶有張量列表的 fit1 花了 0.05196571350097656 秒
fit0 和 dataset0 用了 10.46957802772522 秒
fit0 和 dataset1 用了 7.822799205780029 秒
fit1 with tensorlist 耗時 4.650130748748779 秒

使用@tf.function：
第一次 fit0 與 dataset0 花了 1.4042332172393799 秒
第一次 fit0 與 dataset1 花費了 0.46071624755859375 秒
第一個帶有張量列表的 fit1 耗時 7.3524699211120605 秒
fit0 和 dataset0 用了 15.077088832855225 秒
fit0 和 dataset1 用了 9.136569738388062 秒
fit1 with tensorlist 耗時 2.1366817951202393 秒

第二個是 1 批 100000 個樣本。

輸入大小 = 100000
批量大小 = 100000

沒有@tf.function：
第一次 fit0 與 dataset0 花了 1.1792669296264648 秒
第一次 fit0 與 dataset1 花費了 0.027983427047729492 秒
第一個帶有張量列表的 fit1 花了 0.020987749099731445 秒
fit0 和 dataset0 用了 28.71895956993103 秒
fit0 和 dataset1 用了 2.730872869491577 秒
fit1 with tensorlist 耗時 2.194814682006836 秒

使用@tf.function：
第一次 fit0 與 dataset0 花費了 1.5979444980621338 秒
第一次 fit0 與 dataset1 花費了 0.4557182788848877 秒
第一個帶有張量列表的 fit1 花了 0.3708038330078125 秒
fit0 和 dataset0 用了 36.43854784965515 秒
fit0 和 dataset1 用了 9.819332122802734 秒
fit1 with tensorlist 耗時 2.1136972904205322 秒

問題：

1. 為什么tf.data.Dataset在與tf.function包裝時提供最差的性能？
2. 即使 dataset0 和 dataset1 在功能上是等效的。 兩者之間的底層區別是什么？ 為什么數據集 1 的性能比數據集 0 好？
3. FIT1與tf.function得到了最好的長期演出。
   • 是否可以使用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)

可以在此處找到更多信息。

我在有和沒有tf.function 的情況下測試了fit0 、 fit1 ，但是通過這些更改，我總是通過使用tf.function獲得更好的性能，因此只會顯示后者。

使用的 input_size 是 10 倍。 測試結果如下：

第一個測試是 1000 個批次，每個批次 100 個樣品。
請注意，與 num_unroll = 1 相比，num_unroll = 5 提高了性能。設置 num_unroll > 5 不會提供任何進一步的改進。

輸入大小 = 100000
批量大小 = 100
num_epochs = 100
num_unroll = 5

第一次 fit0 與 dataset0 花了 2.2224882999999993 秒
第一次 fit0 與 dataset1 花費了 0.804360700000001 秒
第一個帶有張量列表的 fit1 耗時 88.2123332 秒
fit0 和 dataset0 用了 35.27911590000001 秒
fit0 和 dataset1 用了 20.370243099999982 秒
fit1 with tensorlist 耗時 23.66727979999999 秒

第二個是 1 批 1000000 個樣本。

輸入大小 = 100000
批量大小 = 100000

輸入大小 = 1000000
批量大小 = 1000000
num_epochs = 100
num_unroll = 1

第一次 fit0 與 dataset0 花費了 4.3616363 秒
第一次 fit0 與 dataset1 花費了 0.7977632000000003 秒
第一個帶有張量列表的 fit1 花了 0.7329889000000005 秒
fit0 和 dataset0 用了 21.131495899999997 秒
fit0 和 dataset1 用了 19.915148600000002 秒
fit1 with tensorlist 耗時 19.817472700000003 秒

上面的結果可以回答這個問題：

3. FIT1與tf.function得到了最好的長期演出。
   • 是否可以使用tf.data.Dataset實現相同的性能？
   • 為什么初始化需要這么多時間？
     當使用 100 個批次時，第一次運行需要 7.3524699211120605 秒，這個時間隨着批次數量的增加而增加。 我猜是因為簽名正在創建一個更大的圖，展開不同批次的計算。 不過，我看不到任何並行化的機會，因為每個批次都依賴於前一個批次的結果。

通過檢查TensorBoard上圖形的結構，很容易看出在fit1函數上使用簽名會通過完全展開循環創建一個非常大的圖形。 這提供了更好的性能，但創建圖形需要很長時間，並且很可能過度使用內存，這使其無法用於更復雜的問題。
但是，如上所示，使用tf.data.Dataset可以實現相同的性能，只需要幾個展開的循環和隨之而來的圖大小的改進。