CODEWITHSUNDEEP
pythonscikit-learnmnist
Shopping Deals
Shopping Deals

Reputation: 33

Python MNIST dataset using sklearn, select specific digits

I’m training few models on MNIST dataset using Sklearn, how do I train the linear model using only two digits 4 and 9 (two classes) from the MNIST dataset?

Upvotes: 2

Views: 2408

Answers (1)

seralouk
seralouk

Reputation: 33147

So you only want to use the images of the digit 4 and 9.

You need indexing like X[np.logical_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)

To use only the digit 4:

X = digits.data
y = digits.target

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

Upvotes: 5

Related Questions