我正在处理MNIST数据集，并探索数据以进行绘图，但在尝试从数据集中提取不同类别时遇到了问题。

import matplotlib.pyplot as pltimport numpy as npfrom tensorflow.examples.tutorials.mnist import input_datamnist = input_data.read_data_sets("MNIST_data/", one_hot=False)dataWithLabels = zip(mnist.train.labels, mnist.train.images)digitDict = {}for i in range(0,10):    digitDict[i] = []for i in dataWithLabels:    digitDict[i[0][0]].append(i[1])for i in range(0,10):    digitDict[i] = np.matrix(digitDict[i])    print("Digit {0} matrix shape: {1}".format(i,digitDict[i].shape))

输出结果是：

Digit 0 matrix shape: (49556, 784)Digit 1 matrix shape: (5444, 784)Digit 2 matrix shape: (1, 0)Digit 3 matrix shape: (1, 0)Digit 4 matrix shape: (1, 0)Digit 5 matrix shape: (1, 0)Digit 6 matrix shape: (1, 0)Digit 7 matrix shape: (1, 0)Digit 8 matrix shape: (1, 0)Digit 9 matrix shape: (1, 0)

但应该是：

Digit 0 matrix shape: (5444, 784)Digit 1 matrix shape: (6179, 784)Digit 2 matrix shape: (5470, 784)Digit 3 matrix shape: (5638, 784)Digit 4 matrix shape: (5307, 784)Digit 5 matrix shape: (4987, 784)Digit 6 matrix shape: (5417, 784)Digit 7 matrix shape: (5715, 784)Digit 8 matrix shape: (5389, 784)Digit 9 matrix shape: (5454, 784)

回答：

我无法重现你的输出（代码对我来说会报错），但你可以用更简单的方法实现：

import numpy as npfrom tensorflow.examples.tutorials.mnist import input_datamnist = input_data.read_data_sets("MNIST_data/", one_hot=False)digitDict = {}for i in range(10):    mask = (mnist.train.labels == i)    digitDict[i] = mnist.train.images[mask]for i in range(10):    print("Digit {0} matrix shape: {1}".format(i,digitDict[i].shape))

这种方法比你的版本要高效得多。代码mnist.train.labels == i会生成一个布尔数组，显示数组中哪些条目等于i，然后我们使用这个布尔数组来获取相应的图像。