将 tf.data.Dataset 包装到 tf.function 中会提高性能吗？

Question

鉴于以下两个示例，在签名时是否有性能改进tf.data.Dataset？

数据集不在 tf.function 中

import tensorflow as tf


class MyModel(tf.keras.Model):

    def call(self, inputs):
        return tf.ones([1, 1]) * inputs


model = MyModel()
model2 = MyModel()


@tf.function
def train_step(data):
    output = model(data)
    output = model2(output)
    return output


dataset = tf.data.Dataset.from_tensors(tf.ones([1, 1]))

for data in dataset:
    train_step(data)

tf.function 中的数据集

import tensorflow as tf


class MyModel(tf.keras.Model):

    def call(self, inputs):
        return tf.ones([1, 1]) * inputs


model = MyModel()
model2 = MyModel()


@tf.function
def train():
    dataset = tf.data.Dataset.from_tensors(tf.ones([1, 1]))
    def train_step(data):
        output = model(data)
        output = model2(output)
        return output
    for data in dataset:
        train_step(data)


train()

Answer 1

添加@tf.function确实会显着提高速度。看看这个：

import tensorflow as tf

data = tf.random.normal((1000, 10, 10, 1))
dataset = tf.data.Dataset.from_tensors(data).batch(10)

def iterate_1(dataset):
    for x in dataset:
        x = x

@tf.function
def iterate_2(dataset):
    for x in dataset:
        x = x

%timeit -n 1000 iterate_1(dataset) # 1.46 ms ± 8.2 µs per loop
%timeit -n 1000 iterate_2(dataset) # 239 µs ± 10.2 µs per loop

如您所见，迭代@tf.function速度提高了 6 倍以上。