使用Pandas循环调用Keras模型时出现KeyError的问题排查

问题背景 #

我创建了3个Keras模型（对应3种不同岩性数据），要在大型数据集上完成批量预测：通过df_cut的Rock列选择对应模型，用同长度的df_big_X数据生成预测值，填充到df_cut的Rv_synth列。

数据样例 #

df_cut样例 #

DEPT    RXOI    SSW TOPS    AT_90   RH39_1DF    RV39_1DF    Rock    GR  ZONES   Rv_synth
68652   3432.60 13.239096   19608.101563    10.0    12.784300   11.323830   29.956650   0.000000    121.453171  1   0
68653   3432.65 28.610432   19782.144531    10.0    14.439166   12.498925   31.454845   0.000000    114.461037  1   0
68654   3432.70 15.997140   19726.130859    10.0    16.094030   13.674020   32.953041   0.000000    107.468910  1   0
68655   3432.75 13.241823   19628.154297    10.0    18.447453   15.684806   35.002056   0.000000    100.766411  1   0
68656   3432.80 12.579830   20877.529297    10.0    20.800871   17.695589   37.051071   0.000000    94.063919   1   0

df_big_X样例 #

RXOI   AT_90   GR
68652   13.239096   12.784300   121.453171
68653   28.610432   14.439166   114.461037
68654   15.997140   16.094030   107.468910
68655   13.241823   18.447453   100.766411
68656   12.579830   20.800871   94.063919
...

错误代码与报错信息 #

运行的循环代码 #

for i in range(0, len(df_big_X)):
    if df_cut['Rock'][i].array == 1:
        y_pred = model_sand_n.predict(df_big_X[i:i+1])
        df_cut['Rv_synth'][i] = y_pred
    elif df_cut['Rock'][i] == 0:
        y_pred = model_shale_n.predict(df_big_X[i:i+1])
        df_cut['Rv_synth'][i] = y_pred
    elif df_cut['Rock'][i] == 6:
         y_pred = model_cal_n.predict(df_big_X[i:i+1])
         df_cut['Rv_synth'][i] = y_pred       

报错信息 #

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File c:\Users\makas\anaconda3\Lib\site-packages\pandas\core\indexes\base.py:3791, in Index.get_loc(self, key)
   3790 try:
-> 3791     return self._engine.get_loc(casted_key)
   3792 except KeyError as err:

File index.pyx:152, in pandas._libs.index.IndexEngine.get_loc()

File index.pyx:181, in pandas._libs.index.IndexEngine.get_loc()

File pandas\_libs\hashtable_class_helper.pxi:2606, in pandas._libs.hashtable.Int64HashTable.get_item()

File pandas\_libs\hashtable_class_helper.pxi:2630, in pandas._libs.hashtable.Int64HashTable.get_item()

KeyError: 0

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Cell In[60], line 2
      1 for i in range(0, len(df_big_X)):
----> 2     if df_cut['Rock'][i].array == 1:
      3         y_pred = model_sand_n.predict(df_big_X[i:i+1])
      4         df_cut['Rv_synth'][i] = y_pred
...
   3801     #  InvalidIndexError. Otherwise we fall through and re-raise
   3802     #  the TypeError.
   3803     self._check_indexing_error(key)

KeyError: 0
Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...

解决方案 #

核心问题分析 #

你的df_cut和df_big_X的索引不是从0开始的连续整数序列（样例里索引是68652、68653这类值），用循环变量i直接通过df_cut['Rock'][i]访问时，pandas会把i当作索引值去查找，而你的索引里没有0，所以触发KeyError。另外第一个条件里多写了.array，这会导致取到的是数组对象而非数值，也是错误的。

修正后的代码 #

改用iloc进行位置索引，确保按行的物理位置而非索引值访问：

for i in range(len(df_big_X)):
    # 用iloc按位置取Rock列的值
    rock_type = df_cut['Rock'].iloc[i]
    # 按位置取对应的输入数据
    x_input = df_big_X.iloc[i:i+1]
    
    if rock_type == 1:
        y_pred = model_sand_n.predict(x_input, verbose=0)
    elif rock_type == 0:
        y_pred = model_shale_n.predict(x_input, verbose=0)
    elif rock_type == 6:
        y_pred = model_cal_n.predict(x_input, verbose=0)
    else:
        # 处理未匹配的岩性类型，比如设为NaN或0
        y_pred = [0]
    
    # 按位置赋值
    df_cut['Rv_synth'].iloc[i] = y_pred[0]

性能优化方案（针对大数据集） #

循环逐行预测效率较低，建议按Rock分组批量预测：

# 为df_big_X添加Rock列，和df_cut对齐
df_big_X['Rock'] = df_cut['Rock'].values

# 批量处理每种岩性
for rock_val, group in df_big_X.groupby('Rock'):
    if rock_val == 1:
        preds = model_sand_n.predict(group.drop('Rock', axis=1), verbose=0)
    elif rock_val == 0:
        preds = model_shale_n.predict(group.drop('Rock', axis=1), verbose=0)
    elif rock_val == 6:
        preds = model_cal_n.predict(group.drop('Rock', axis=1), verbose=0)
    else:
        preds = [0]*len(group)
    
    # 按索引赋值
    df_cut.loc[group.index, 'Rv_synth'] = preds

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