【Tensorflow教程笔记】常用模块 tf.function ：图执行模式

最新推荐文章于 2024-03-05 17:03:28 发布

_APTX4869

最新推荐文章于 2024-03-05 17:03:28 发布

阅读量1.5w

点赞数 10

分类专栏： Tensorflow

本文链接：https://blog.csdn.net/nanke_4869/article/details/114220354

版权

本文详细介绍了 TensorFlow 中的 `tf.function` 模块，用于将模型转换为图执行模式，提高性能和便于部署。内容涵盖 `tf.function` 的基础使用方法、内在机制、AutoGraph 的工作原理以及如何在图执行模式下使用 `tf.Session`。通过示例展示了 `tf.function` 如何优化代码执行效率，并解释了其在处理 Python 控制流和数据结构时的特性和注意事项。

摘要由CSDN通过智能技术生成

虽然默认的即时执行模式（Eager Execution）为我们带来了灵活及易调试的特性，但在特定的场合，例如追求高性能或部署模型时，我们依然希望使用 TensorFlow 1.X 中默认的图执行模式（Graph Execution），将模型转换为高效的 TensorFlow 图模型。此时，TensorFlow 2 为我们提供了 tf.function 模块，结合 AutoGraph 机制，使得我们仅需加入一个简单的 @tf.function 修饰符，就能轻松将模型以图执行模式运行。

`tf.function` 基础使用方法

在 TensorFlow 2 中，推荐使用 tf.function （而非 1.X 中的 tf.Session ）实现图执行模式，从而将模型转换为易于部署且高性能的 TensorFlow 图模型。只需要将我们希望以图执行模式运行的代码封装在一个函数内，并在函数前加上 @tf.function 即可，如下例所示。关于图执行模式的深入探讨可参考附录。

并不是任何函数都可以被 @tf.function 修饰！@tf.function 使用静态编译将函数内的代码转换成计算图，因此对函数内可使用的语句有一定限制（仅支持 Python 语言的一个子集），且需要函数内的操作本身能够被构建为计算图。建议在函数内只使用 TensorFlow 的原生操作，不要使用过于复杂的 Python 语句，函数参数只包括 TensorFlow 张量或 NumPy 数组，并最好是能够按照计算图的思想去构建函数（换言之，@tf.function 只是给了你一种更方便的写计算图的方法，而不是一颗能给任何函数加速的银色子弹）。

import tensorflow as tf
import time
from zh.model.mnist.cnn import CNN
from zh.model.utils import MNISTLoader

num_batches = 1000
batch_size = 50
learning_rate = 0.001
data_loader = MNISTLoader()
model = CNN()
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

@tf.function
def train_one_step(X, y):    
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        # 注意这里使用了TensorFlow内置的tf.print()。@tf.function不支持Python内置的print方法
        tf.print("loss", loss)
    grads = tape.gradient(loss, model.variables)    
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))

start