TensorFlow构建大尺寸视频tf.Dataset避免内核崩溃的方法

问题背景 #

• 任务目标：使用TensorFlow搭建视频回归模型，实现基于图像预测火灾释放热量的功能。前期CNN方案已取得良好效果，当前尝试实现长时循环卷积网络。
• 数据参数：单条视频输入维度为(161,540,960,3)，即单视频包含161帧，单帧尺寸为540*960、3通道；共193条视频样本，全量数据张量维度为(193,161,540,960,3)。
• 运行环境：Google Colab Pro
• 故障现象：原有Dataset构建代码最多处理2个视频文件就会触发内核崩溃，需要更简洁、内存效率更高的tf.Dataset实现方案。

原有故障代码 #

def create_dataset(df,num_experiments,max_len):
  for i in range(num_experiments):
    for j in range(max_len):
      if j == 0:
        gotten_image = get_img(df.filename_list[i][j]) #Initialize image
        images = gotten_image

        mask = tf.constant([1]) #Initialize mask
        labels = tf.constant([df.hrr_list[i][j]]) #Initialize labels
      else:
        try:
          gotten_image = get_img(df.filename_list[i][j]) #Append image if it exists
          
          mask = tf.concat([mask, tf.constant([1])],0) #Append mask if image exists
          labels = tf.concat([labels, tf.constant([df.hrr_list[i][j]])],0) #Append label if it exists

        except:
          gotten_image = tf.zeros([1,540,960,3]) #Append zeroes if no image exists
          
          mask = tf.concat([mask, tf.constant([0])],0)  #Append zeroes if no image exists
          labels = tf.concat([labels, tf.constant([0.])],0) #Append zeroes if no image exists


        images = tf.concat([images,gotten_image],0) #Layer images from a video one after the other
    
    if i == 0:
      full_images = images
      full_images = tf.expand_dims(full_images,0) #Initialize list of videos

      full_mask = tf.expand_dims(mask,0)
      full_labels = tf.expand_dims(labels,0)
    else:
      images = tf.expand_dims(images,0) 
      full_images = tf.concat([full_images,images],0) #Append one video to the rest

      mask = tf.expand_dims(mask,0)
      labels = tf.expand_dims(labels,0)
      full_mask = tf.concat([full_mask,mask],0) #Append corresponding mask
      full_labels = tf.concat([full_labels,labels],0) #Append coressponding labels


  dataset = tf.data.Dataset.from_tensor_slices(((full_images, full_mask), full_labels)) #Create dataset

  return dataset

故障根因 #

原有实现存在两个致命的内存问题：

1. 循环中反复调用tf.concat拼接张量，每次拼接都会生成全新的张量副本，内存占用随样本量呈平方级增长。
2. 构建Dataset前就将全部193条视频的所有帧解码加载到内存，按float32精度计算仅图像数据就需要超过170GB内存，远超Colab Pro的内存上限，必然触发内核崩溃。

优化实现方案 #

核心思路是放弃全量预加载逻辑，利用tf.data的懒加载特性实现流式读取，仅在训练需要时加载对应批次的数据，内存占用完全可控。

import tensorflow as tf

def create_dataset(df, num_experiments, max_len, batch_size=2, shuffle=True):
    # 第一步：仅构造轻量元数据列表，不加载实际图像，内存占用不足10MB
    meta_list = []
    for exp_id in range(num_experiments):
        frame_paths = df.filename_list[exp_id]
        frame_hrrs = df.hrr_list[exp_id]
        # 提前将路径、标签、掩码补齐到固定长度
        paths_pad, labels_pad, mask_pad = [], [], []
        valid_frame_num = min(len(frame_paths), max_len)
        for frame_id in range(max_len):
            if frame_id < valid_frame_num:
                paths_pad.append(frame_paths[frame_id])
                labels_pad.append(float(frame_hrrs[frame_id]))
                mask_pad.append(1)
            else:
                paths_pad.append("") # 缺失帧用空路径占位
                labels_pad.append(0.0)
                mask_pad.append(0)
        meta_list.append((paths_pad, labels_pad, mask_pad))
    
    # 从元数据生成器创建Dataset，初始化阶段无图像加载开销
    def data_generator():
        for item in meta_list:
            yield item
    
    dataset = tf.data.Dataset.from_generator(
        data_generator,
        output_signature=(
            tf.TensorSpec(shape=(max_len,), dtype=tf.string),
            tf.TensorSpec(shape=(max_len,), dtype=tf.float32),
            tf.TensorSpec(shape=(max_len,), dtype=tf.int32)
        )
    )

    # 定义单视频加载逻辑，逐帧解码
    def load_video_sample(paths, labels, mask):
        frames = []
        for idx in tf.range(max_len):
            frame_path = paths[idx]
            if frame_path == tf.constant(""):
                # 缺失帧返回零张量
                frame = tf.zeros((540, 960, 3), dtype=tf.float32)
            else:
                # 调用原有图像读取函数
                frame = tf.numpy_function(
                    func=get_img,
                    inp=[frame_path],
                    Tout=tf.float32
                )
                frame.set_shape((540, 960, 3))
            frames.append(frame)
        # 仅堆叠单视频的帧，内存开销可控
        frames = tf.stack(frames, axis=0)
        return (frames, mask), labels

    # 并行加载数据，自动调整并行度
    dataset = dataset.map(load_video_sample, num_parallel_calls=tf.data.AUTOTUNE)
    if shuffle:
        # 小缓冲区做乱序，避免大缓冲占用过多内存
        dataset = dataset.shuffle(buffer_size=8)
    dataset = dataset.batch(batch_size)
    # 开启预取，让数据加载和模型计算重叠，提升训练速度
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

使用说明 #

• 默认batch_size=2时，同时驻留内存的视频仅2个，图像数据内存占用约1.8GB，Colab Pro环境可稳定运行；如果内存仍有富余可适当调大batch_size。
• 如果需要进一步降低内存占用，可在读取单帧后增加resize操作，将帧分辨率下调到模型实际需要的输入尺寸。
• 该实现完全兼容模型训练的输入要求，返回的Dataset格式和原有实现一致，无需修改模型侧代码。

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