I am just starting out with tensorflow and I want to test something only on the 0's and 1's from the MNIST images. Is there a way to import only these images?
Assuming you are using
from tensorflow.examples.tutorials.mnist import input_data
No, there is no function or argument in that file... What you can do is load all data, and select only the ones and zeros.
(x_train, y_train), (x_test, y_test) = mnist.load_data()
train_mask = np.isin(y_train, [0, 1])
x = x_train[train_mask]
y = y_train[train_mask]