设置batch_size=1仍报Dst tensor未初始化错误的技术求助

可能原因 #

• 全数据集内存预分配问题：TensorFlow通过tf.data.Dataset.from_tensor_slices加载完整大张量时，会尝试将整个数据集一次性加载到内存并完成张量初始化，哪怕设置batch_size=1也会触发这个预分配流程。而截取前20个样本时，张量体积小，初始化过程不会触发底层内存分配错误。
• Jupyter内存残留：Jupyter内核会保留之前运行产生的内存占用，即使重新执行代码，未完全释放的旧大张量会挤占内存空间，导致新的完整数据集张量初始化失败，小数据集因内存需求低可正常运行。
• TensorFlow 2.10版本bug：该版本针对大维度张量（如你的131072维度）的初始化可能存在底层实现问题，小张量未触发该bug。
• 数据集潜在损坏：完整数据集中可能存在个别样本数据异常（如NaN、数据类型不匹配），导致整体张量初始化失败，而前20个样本无异常。

解决方法 #

• 改用流式加载方式：避免一次性加载全量数据，使用生成器流式读取样本，比如：

  def data_generator():
      for idx in range(len(train_x)):
          yield train_x[idx], train_y[idx]
  
  dataset = tf.data.Dataset.from_generator(
      data_generator,
      output_signature=(
          tf.TensorSpec(shape=(131072, 2), dtype=tf.float32),
          tf.TensorSpec(shape=YOUR_LABEL_SHAPE, dtype=tf.float32)
      )
  ).batch(1)
  

• 清理Jupyter内存：重启内核彻底释放残留内存，或在代码开头手动清理：

  import gc
  # 删除旧的大变量并强制垃圾回收
  for var_name in ['train_x', 'train_y']:
      if var_name in locals():
          del locals()[var_name]
  gc.collect()
  

• 调整TF内存分配策略（GPU环境）：设置显存按需分配，避免预占全部显存：

  import tensorflow as tf
  gpus = tf.config.list_physical_devices('GPU')
  if gpus:
      for gpu in gpus:
          tf.config.experimental.set_memory_growth(gpu, True)
  

• 校验数据集完整性：检查全量数据中是否存在异常值：

  import numpy as np
  # 检查是否有NaN或Inf值
  print(np.isnan(train_x).any())
  print(np.isinf(train_x).any())
  # 检查数据类型是否统一
  print(train_x.dtype)
  

• 更换TensorFlow版本：尝试升级到2.11及以上版本，或降级到2.9版本，规避版本bug。

内容的提问来源于stack exchange，提问作者Mitchell van Zuylen