使用sklearn的Python MNIST数据集，选择特定数字

Question

我正在使用 Sklearn 在 MNIST 数据集上训练几个模型，如何仅使用 MNIST 数据集中的两位数字 4 和 9(两个类别)来训练线性模型？

• 如何选择我的X_test,X_train, y_test,y_train？

Answer 1

所以你只想使用数字 4 和 9 的图像。

您需要像X[np.逻辑_or(y == 4, y == 9)]这样的索引:

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_digits

digits = load_digits()

X = digits.data
y = digits.target

#Select only the digit 4 and 9 images
X = X[np.logical_or(y == 4, y == 9)]
y = y[np.logical_or(y == 4, y == 9)]

# verify selection
np.unique(y)
#array([4, 9])

# Now split them
X_train, X_test, y_train, y_test = train_test_split(
    X, y, train_size=200, test_size=100)

<小时/>

仅使用数字 4:

X = digits.data
y = digits.target

#Select only the digit 4 and 9 images
X = X[y == 4]
y = y[y == 4]