TensorFlow CSV Dataset读取时GPU利用率极低问题咨询

核心原因是数据输入管道的速度跟不上GPU的计算需求，导致GPU大部分时间处于等待状态。下面是针对性的优化措施：

1. 将CSV转换为TFRecord格式 #

CSV是文本格式，解析开销大，TFRecord作为TensorFlow原生二进制格式，读取和解析效率远高于CSV。先批量转换所有CSV文件为TFRecord，再从TFRecord读取数据：

# 示例：将CSV转换为TFRecord
def serialize_example(feature_dict):
    feature = {}
    for col in columns:
        if col == "target":
            feature[col] = tf.train.Feature(int64_list=tf.train.Int64List(value=[feature_dict[col]]))
        else:
            # 根据column_types调整对应特征类型
            feature[col] = tf.train.Feature(float_list=tf.train.FloatList(value=[feature_dict[col]]))
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

# 读取CSV并写入TFRecord
raw_ds = tf.data.experimental.make_csv_dataset(
    file_pattern="gs://folder/*.csv",
    batch_size=1000,
    column_names=columns,
    column_defaults=column_types,
    num_epochs=1,
    ignore_errors=True
)

writer = tf.io.TFRecordWriter("gs://folder/data.tfrecord")
for batch in raw_ds:
    for i in range(batch["target"].shape[0]):
        feature_dict = {k: batch[k][i].numpy() for k in columns + ["target"]}
        serialized_example = serialize_example(feature_dict)
        writer.write(serialized_example)
writer.close()

# 从TFRecord读取数据
def parse_example(example_proto):
    feature_desc = {}
    for col in columns:
        if col == "target":
            feature_desc[col] = tf.io.FixedLenFeature([], tf.int64)
        else:
            feature_desc[col] = tf.io.FixedLenFeature([], tf.float32)
    parsed = tf.io.parse_single_example(example_proto, feature_desc)
    return {k: parsed[k] for k in columns}, parsed["target"]

my_ds = tf.data.TFRecordDataset("gs://folder/data.tfrecord", num_parallel_reads=tf.data.AUTOTUNE)
my_ds = my_ds.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)

2. 优化tf.data管道性能 #

• 添加预取：在数据集最后加上prefetch(tf.data.AUTOTUNE)，让数据准备与GPU计算并行，避免GPU等待数据。
• 启用缓存：如果数据集能放入本地内存，添加cache()；如果数据量过大，使用tf.data.experimental.cache("/path/to/local/cache")将缓存写入本地高速存储。
• 并行预处理：如果有自定义预处理步骤，使用map(..., num_parallel_calls=tf.data.AUTOTUNE)并行执行逻辑。
• 自动调整读取并行度：将num_parallel_reads设置为tf.data.AUTOTUNE，让TensorFlow根据系统资源自动调整读取并行数。

3. 调整分布式训练配置 #

• 调整全局批量大小：MirroredStrategy下，batch_size是全局批量，会自动均分至每个GPU（比如全局批量64，双GPU各处理32）。适当增大全局批量（如128、256），让GPU有足够的数据处理。
• 确保模型与数据集在策略作用域内：

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # 构建并编译模型
    model = build_your_model()
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")

# 数据集批量与预取操作放在策略作用域后，确保分布式兼容
my_ds = my_ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)

4. GCS存储优化 #

• 区域匹配：确保GCS存储桶与GPU所在计算区域一致，减少跨区域网络延迟。
• 本地缓存GCS文件：设置GOOGLE_CLOUD_CACHE_DIR环境变量，让TensorFlow将GCS文件缓存到本地磁盘，避免重复下载相同数据。

5. 排查预处理瓶颈 #

使用TensorFlow的Profile工具（TensorBoard中开启Profile），分析数据输入管道各阶段耗时，定位是否存在预处理函数的性能瓶颈，优先用TensorFlow原生操作替代Python逻辑。

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