使用sklearn获取MNIST数据时X[0]触发KeyError:0错误的原因

问题场景 #

我正在跟随《Hands-On Machine Learning with Scikit-Learn and TensorFlow》第3章进行MNIST数据分类任务，在Jupyter Notebook中执行以下代码：

from sklearn.datasets import fetch_openml
mnist = fetch_openml('mnist_784', version=1)
mnist.keys()
# 输出: dict_keys(['data', 'target', 'feature_names', 'DESCR', 'details', 'categories', 'url'])

X, y = mnist["data"], mnist["target"]
X.shape
# 输出: (70000, 784)
y.shape
# 输出: (70000,)

但执行some_digit = X[0]时触发KeyError:0错误，错误信息如下：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/anaconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3360             try:
-> 3361                 return self._engine.get_loc(casted_key)
   3362             except KeyError as err:

~/anaconda3/lib/python3.7/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

~/anaconda3/lib/python3.7/site-packages/pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 0

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

KeyError                                  Traceback (most recent call last)
<ipython-input-43-348a6e96ae02> in <module>
----> 1 some_digit = X[0]

~/anaconda3/lib/python3.7/site-packages/pandas/core/frame.py in __getitem__(self, key)
   3456             if self.columns.nlevels > 1:
   3457                 return self._getitem_multilevel(key)
-> 3458             indexer = self.columns.get_loc(key)
   3459             if is_integer(indexer):
   3460                 indexer = [indexer]

~/anaconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
   3361                 return self._engine.get_loc(casted_key)
   3362             except KeyError as err:
-> 3363                 raise KeyError(key) from err
   3364 
   3365         if is_scalar(key) and isna(key) and not self.hasnans:

KeyError: 0

问题根源 #

• 当前环境中，fetch_openml返回的mnist["data"]是Pandas DataFrame对象，而非书籍预期的NumPy数组。
• 在Pandas DataFrame中，X[0]的语义是查找列名为0的列，但MNIST数据的列名是pixel0到pixel783，不存在该列，因此触发KeyError。
• 旧版scikit-learn中fetch_openml默认返回NumPy数组，此时X[0]会直接取第0行数据，这是书籍中的原始逻辑。

解决方法 #

有两种可行方案：

1. 使用Pandas的行索引方法获取第0行：

some_digit = X.iloc[0]

2. 将DataFrame转换为NumPy数组，还原书籍中的操作逻辑：

X = mnist["data"].values
some_digit = X[0]

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