tf.data.Dataset.from_generator训练时随机出现InvalidArgumentError

问题定位 #

报错核心是InvalidArgumentError: Index out of range using input dim 1; input has only 1 dims，精准定位到生成器中x_start = self.indices[self.counter]这一行。单独遍历数据集或直接用生成器训练正常，但开启prefetch()后随机触发错误，根源是tf.data的多线程并行机制会并发修改生成器的共享状态变量self.counter，导致索引超出self.indices的范围。

解决方案 #

1. 替换共享状态变量，使用局部迭代器 #

将self.indices转为numpy数组，在生成器内部创建局部迭代器，彻底避免多线程状态冲突：

def generate(self):
    # 将Tensor转为numpy数组，规避Tensor操作的线程安全风险
    indices_np = self.indices.numpy()
    while True:
        # 每个epoch重新打乱索引，保证随机性
        np.random.shuffle(indices_np)
        # 按batch_size分组遍历索引
        for batch_start in range(0, len(indices_np), self.batch_size):
            batch_indices = indices_np[batch_start:batch_start+self.batch_size]
            batch_demand = []
            batch_supply = []
            batch_scale_params = []
            batch_targets = []
            
            for idx in batch_indices:
                x_start = idx
                # 原样本提取、预处理逻辑（保持不变）
                demand = self.x_data[
                    x_start - self.scale_window : x_start + self.sequence_length, :, :2
                ]
                supply = self.x_data[
                    x_start - self.scale_window : x_start + self.sequence_length, :, 2:
                ]
                # ...后续多项式特征、归一化等处理...
                
                # 收集处理后的样本
                batch_demand.append(tf.constant(demand, self.x_dtype))
                batch_supply.append(tf.constant(supply, self.x_dtype))
                batch_scale_params.append(tf.constant([self.mmx.data_min_[0], self.mmx.data_max_[0]], self.y_dtype))
                batch_targets.append(tf.constant(
                    self.x_data[
                        x_start + self.sequence_length : x_start + self.sequence_length + self.n_forecast,
                        0,
                        [0, 2],
                    ],
                    self.y_dtype,
                ))
            
            # 组装批次并返回
            yield (
                tf.stack(batch_demand),
                tf.stack(batch_supply),
                tf.stack(batch_scale_params),
                tf.stack(batch_targets)
            )

2. 改用tf.data原生API（推荐） #

抛弃Python生成器，直接用tf.data链式API处理，天然支持并行且无状态冲突：

def create_dataset(self):
    # 将原始数据转为Tensor
    x_data_tensor = tf.convert_to_tensor(self.x_data, dtype=self.x_dtype)
    # 创建索引数据集
    dataset = tf.data.Dataset.from_tensor_slices(self.indices)
    dataset = dataset.shuffle(buffer_size=len(self.indices), seed=1)
    dataset = dataset.batch(self.batch_size)
    
    # 定义批量预处理函数（需将sklearn操作转为tf兼容实现）
    @tf.function
    def preprocess_batch(indices_batch):
        batch_demand = []
        batch_supply = []
        batch_scale_params = []
        batch_targets = []
        
        for idx in indices_batch:
            # 提取序列
            demand = x_data_tensor[idx - self.scale_window : idx + self.sequence_length, :, :2]
            supply = x_data_tensor[idx - self.scale_window : idx + self.sequence_length, :, 2:]
            
            # 多项式特征处理（替换sklearn.PolynomialFeatures）
            poly_layer = tf.keras.layers.PolynomialFeatures(degree=self.poly.degree, include_bias=False)
            demand = tf.reshape(poly_layer(tf.reshape(demand, (-1, 2))), (*demand.shape[:-1], -1))
            supply = tf.reshape(poly_layer(tf.reshape(supply, (-1, 2))), (*supply.shape[:-1], -1))
            
            # 归一化处理（替换sklearn.MinMaxScaler，用tf实现滑动窗口归一）
            demand = tf.reverse(demand, axis=[0])
            supply = tf.reverse(supply, axis=[0])
            # ...实现滑动窗口归一化逻辑...
            
            # 提取目标
            forecast_start = idx + self.sequence_length
            target = x_data_tensor[forecast_start : forecast_start + self.n_forecast, 0, [0, 2]]
            
            batch_demand.append(demand)
            batch_supply.append(supply)
            batch_scale_params.append(scale_param)
            batch_targets.append(target)
        
        return (
            tf.stack(batch_demand),
            tf.stack(batch_supply),
            tf.stack(batch_scale_params),
            tf.stack(batch_targets)
        )
    
    # 启用并行处理和预取
    dataset = dataset.map(preprocess_batch, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

3. 给状态变量加线程锁（临时方案） #

如果必须保留生成器，通过线程锁保护self.counter的读写：

import threading

class BatchGenerator:
    def __init__(self, **kwargs):
        # 其他初始化代码...
        self.counter_lock = threading.Lock()  # 新增锁
    
    def generate(self):
        while True:
            batch_demand = tf.TensorArray(self.x_dtype, self.batch_size)
            # ...其他TensorArray初始化...
            
            for i in range(self.batch_size):
                # 用锁保护counter的读写
                with self.counter_lock:
                    x_start = self.indices[self.counter]
                    # 更新counter的逻辑也放在锁内
                    if self.counter + 1 == len(self.indices):
                        self.counter = 0
                    else:
                        self.counter += 1
                
                # 后续样本处理逻辑保持不变...

核心原因 #

tf.data的prefetch()和map()会启动多线程并行拉取/处理数据，而Python生成器的共享状态变量（如self.counter）没有线程安全保护，多个线程同时读写会导致counter的值异常增大，超出self.indices的长度，最终触发索引越界错误。单线程执行时（如直接遍历生成器）不会出现该问题。

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