如何禁用tf.data自动调优？解决TensorFlow训练卡顿问题

问题描述 #

我当前运行一个源自TensorFlow官方教程的小型模型，希望禁用tf.data的自动调优。训练时出现如下日志信息，程序执行陷入卡顿，怀疑是tf.data的自动调优出现异常。使用环境为TensorFlow v2.4版本，GPU为RTX 2080Ti。

2022-09-14 02:57:26.254717: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 02:57:26.254980: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 02:57:26.255059: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 20480 ms.
2022-09-14 02:57:46.735341: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 02:57:46.735523: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 02:57:46.735562: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 40960 ms.
2022-09-14 02:58:27.695827: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 02:58:27.696006: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 02:58:27.696046: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 02:59:27.696275: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 02:59:27.696469: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 02:59:27.696507: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:00:27.696740: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:00:27.696911: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:00:27.696951: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:01:27.697178: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:01:27.697368: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:01:27.697407: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:02:27.697638: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:02:27.697811: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:02:27.697850: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:03:27.698078: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:03:27.698250: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:03:27.698309: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:04:27.698543: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:04:27.698724: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:04:27.698762: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:05:27.698990: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:05:27.699164: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:05:27.699202: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:06:27.699453: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:06:27.699627: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:06:27.699668: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.
2022-09-14 03:07:27.699893: I tensorflow/core/framework/model.cc:1439] Starting optimization of tunable parameters with GradientDescent
2022-09-14 03:07:27.700112: I tensorflow/core/framework/model.cc:1494] Number of tunable parameters: 0
2022-09-14 03:07:27.700175: I tensorflow/core/kernels/data/model_dataset_op.cc:200] Waiting for 60000 ms.

解决方案 #

在TensorFlow 2.4中，有两种可靠的方式关闭tf.data自动调优：

1. 针对Dataset对象单独设置 #

创建tf.data.Options()实例，禁用自动调优后绑定到你的数据集上：

import tensorflow as tf

# 创建选项并禁用自动调优
options = tf.data.Options()
options.experimental_optimization.autotune.enabled = False

# 将选项应用到数据集
dataset = dataset.with_options(options)

2. 全局禁用（通过环境变量） #

在启动训练脚本前设置环境变量，全局关闭tf.data自动调优：

# Linux/macOS终端执行
export TF_DATA_AUTOTUNE_DISABLE=1

# Windows命令行执行
set TF_DATA_AUTOTUNE_DISABLE=1

从日志可以看到，自动调优模块检测到“可调整参数数量为0”却仍在等待调优完成，这属于异常行为，关闭自动调优后应该能解决程序卡顿问题。

内容的提问来源于stack exchange，提问作者Jueon Park