TensorFlow K折交叉验证遇InternalError：目标张量未初始化

问题概述 #

基于Xception构建手部X光图像的月龄估计回归模型，采用5折交叉验证计算平均绝对误差（MAE）。首次交叉验证运行正常，但循环执行model.predict()时触发GPU内存相关的InternalError，即使将model.predict()的batch_size调至2仍无法解决。数据集包含8000张手部X光图像。

错误详情 #

错误日志 #

640/640 [==============================] - 86s 114ms/step - loss: 0.3346 - mae_months: 17.8703

---------------------------------------------------------------------------
InternalError                             Traceback (most recent call last)
Input In [15], in <cell line: 3>()
      3 for train_index, val_index in kf.split(X_train):
      4     model.fit(X_train[train_index], y_train[train_index], batch_size = 10)
----> 5     pred = model.predict(X_train[val_index], batch_size = 2)
      6     err = mean_absolute_error(y_train[val_index], pred)
      7     cv_mae.append(err)

File ~\anaconda3\lib\site-packages\keras\wrappers\scikit_learn.py:364, in KerasRegressor.predict(self, x, **kwargs)
    350 """Returns predictions for the given test data.
    351 
    352 Args:
   (...)
    361         Predictions.
    362 """
    363 kwargs = self.filter_sk_params(Sequential.predict, kwargs)
--> 364 return np.squeeze(self.model.predict(x, **kwargs))

File ~\anaconda3\lib\site-packages\keras\utils\traceback_utils.py:67, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     65 except Exception as e:  # pylint: disable=broad-except
     66   filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67   raise e.with_traceback(filtered_tb) from None
     68 finally:
     69   del filtered_tb

File ~\anaconda3\lib\site-packages\tensorflow\python\framework\constant_op.py:102, in convert_to_eager_tensor(value, ctx, dtype)
    100     dtype = dtypes.as_dtype(dtype).as_datatype_enum
    101 ctx.ensure_initialized()
--> 102 return ops.EagerTensor(value, ctx.device_name, dtype)

InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run _EagerConst: Dst tensor is not initialized.

关键错误提示 #

InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run _EagerConst: Dst tensor is not initialized.

核心原因 #

1. 模型实例重复使用：循环外创建的KerasRegressor实例在每次fold训练后保留权重，GPU内存持续累积占用，后续predict时内存不足无法初始化张量。
2. 全量数据加载：通过next(train_gen)一次性将所有训练图像加载到X_train数组，占用大量CPU和GPU内存，交叉验证的切片操作进一步加剧内存压力。
3. GPU内存未自动释放：TensorFlow默认不会主动释放GPU内存，多次训练迭代后内存耗尽。

解决方案 #

1. 每次交叉验证重新构建模型 #

在循环内部重新初始化模型，避免权重残留和内存累积：

# KFold
n_splits = 5
kf = KFold(n_splits = n_splits, shuffle = True, random_state = 42)

cv_mae = []

for train_index, val_index in kf.split(X_train):
    # 每次fold重新创建模型
    model = KerasRegressor(build_fn = boneage_model)
    model.fit(X_train[train_index], y_train[train_index], batch_size = 16)
    pred = model.predict(X_train[val_index], batch_size = 2)
    err = mean_absolute_error(y_train[val_index], pred)
    cv_mae.append(err)
    
    # 清理GPU内存
    tf.keras.backend.clear_session()
    del model

2. 改用生成器分批加载数据（推荐） #

避免一次性加载全量数据，用生成器处理训练和验证集：

# 修改数据生成逻辑，为每个fold创建对应的生成器
def get_fold_generator(df, indices, img_data_gen, path_col, y_col, target_size, batch_size):
    fold_df = df.iloc[indices].reset_index(drop=True)
    return flow_from_dataframe(img_data_gen, fold_df,
                               path_col=path_col, y_col=y_col,
                               target_size=target_size, color_mode='rgb',
                               batch_size=batch_size, shuffle=False)

# 初始化KFold
n_splits = 5
kf = KFold(n_splits=n_splits, shuffle=True, random_state=42)

cv_mae = []

for train_idx, val_idx in kf.split(train_df):
    # 创建当前fold的训练和验证生成器
    train_fold_gen = get_fold_generator(train_df, train_idx, core_idg, 'path', 'boneage_zscore', IMG_SIZE, 16)
    val_fold_gen = get_fold_generator(train_df, val_idx, core_idg, 'path', 'boneage_zscore', IMG_SIZE, 2)
    
    # 构建并训练模型
    model = boneage_model()
    model.fit(train_fold_gen, epochs=1)
    
    # 获取验证集真实值和预测值
    val_true = []
    val_pred = []
    for x_val, y_val in val_fold_gen:
        val_true.extend(y_val)
        val_pred.extend(model.predict(x_val, batch_size=2).flatten())
    
    # 计算MAE
    err = mean_absolute_error(np.array(val_true)*boneage_div, np.array(val_pred)*boneage_div)
    cv_mae.append(err)
    
    # 清理内存
    tf.keras.backend.clear_session()
    del model

3. 优化GPU内存配置 #

确保GPU内存增长模式开启，避免内存预分配耗尽资源：

physical_devices = tf.config.list_physical_devices('GPU')
if physical_devices:
    try:
        for device in physical_devices:
            tf.config.experimental.set_memory_growth(device, True)
        print("GPU memory growth enabled")
    except RuntimeError as e:
        print(e)

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