使用TensorFlow加载视频帧的最优方法：如何将keras.utils.Sequence替换为tf.data

你遇到的GPU无调用的核心原因是：原来的keras.utils.Sequence逻辑里，图片读取、resize全部在CPU单线程执行，GPU训练完一个批次后需要等待CPU预处理完下一批次数据才能继续，全程GPU空闲率极高，看起来就像没有被调用。用tf.data实现的核心是把数据加载/预处理做成异步并行流水线，让CPU在GPU训练当前批次的时候就提前处理好下一批次数据，彻底打满GPU利用率。

以下是和你原生成器逻辑完全对齐的tf.data实现代码：

import tensorflow as np
import tensorflow as tf

# 替换为你自己的实际参数
IMG_SIZE = 224
BATCH_SIZE = 8
NUM_CLASSES = 16

def load_single_video(frame_paths, label_str):
    # 标签处理逻辑和原代码完全一致：取字符串第3位后转int减1，再做独热编码
    label = tf.strings.substr(label_str, 3, tf.strings.length(label_str)-3)
    label = tf.strings.to_number(label, out_type=tf.int32) - 1
    label = tf.one_hot(label, depth=NUM_CLASSES)

    # 单帧处理逻辑：读取+resize，输出归一化0-1的张量，和skimage.resize效果一致
    def process_single_frame(frame_path):
        img = tf.io.read_file(frame_path)
        img = tf.image.decode_jpeg(img, channels=3)
        img = tf.image.resize(img, (IMG_SIZE, IMG_SIZE))
        return img

    # 堆叠当前视频所有帧，输出格式为 (#frames, IMG_SIZE, IMG_SIZE, 3)
    frames = tf.map_fn(process_single_frame, frame_paths, fn_output_signature=tf.float32)
    return frames, label

完成加载函数定义后，按以下步骤构建数据集即可：

# x_set为你的原始输入列表，每个元素是单个视频的所有帧路径列表
# y_set为和x_set一一对应的标签字符串列表
x_set = ...
y_set = ...

# 构建基础数据集
dataset = tf.data.Dataset.from_tensor_slices((x_set, y_set))
# 开启多线程并行预处理，自动适配CPU核心数
dataset = dataset.map(load_single_video, num_parallel_calls=tf.data.AUTOTUNE)
# 批次化，输出维度和原生成器完全一致：(BATCH_SIZE, #frames, IMG_SIZE, IMG_SIZE, 3)
dataset = dataset.batch(BATCH_SIZE)
# 开启预加载队列，GPU训练当前批次时CPU提前加载好后续批次，消除GPU等待
dataset = dataset.prefetch(tf.data.AUTOTUNE)

训练时直接把数据集传入model.fit即可，不需要再使用原来的Sequence生成器：

model.fit(
    dataset,
    epochs=10
)

特殊情况适配 #

如果你的不同视频帧数不一致，需要把普通batch替换为填充批次，保证同批次张量维度统一：

dataset = dataset.padded_batch(
    BATCH_SIZE,
    padded_shapes=([None, IMG_SIZE, IMG_SIZE, 3], [NUM_CLASSES])
)

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