在上一篇博客 digit_eye开发记录(1): C++读取MNIST数据集 中解读了 IDX 文件格式，并使用 C++ 语言完成了 MNIST 数据集的解析，第6小节给出的完整代码有146行之多。使用 Python 读取则可以省略70%的代码，只用不到50行代码完成相同功能。

读取 buffer

np.frombuffer(buf, dtype, count, offset)

说明：

• buf: buffer，从文件读出来的
• dtype: 从buf读取时，按什么类型读取数据，或者说，读取的基本单位是什么
• count: 从buf读取时，读取多少个基本单位
• offset: 从buf读取时，指针首先偏移多少个字节

读取 magic number

magic number 是 mnist 文件的前4个字节。 以二进制形式打开后，读取4字节即可：

import numpy as npwith open(filename, 'rb') as fin:buf = bytearray(fin.read())
magic = np.frombuffer(buf, np.uint8, count=4)
print(magic)

读取维度信息

回忆一下 magic numbers 的构成： 前两个字节是0，第三个字节是类型，第四个字节是维度数量 num_dims。
根据 num_dims 的取值，读取对应数量的字节，得到对应的维度信息。每个维度都是一个 int32 大小。

注意 MSB 到 LSB 的转换，通过 dtype=np.dtype('>u4') 指定， >u4 意思是：以MSB序，读取4个byte.

对于图像数据：

num_dims = magic[3]
dims = np.frombuffer(buf, dtype=np.dtype('>u4'), count=num_dims, offset=4)

对于label数据：

dims = np.frombuffer(buf, dtype=np.dtype('>u4'), count=num_dims, offset=4)
num_labels = dims[0]

读取图像像素

很容易想到使用 OOP 方式，定义 DataSet 类，在成员 self.images 中保存图像；于是乎，很“毛躁”的写出如下糟糕代码：

class DataSet:def __init__(self):self.images = []self.labels = []def load_images(self, filename):...for i in range(num_images):self.images.append(...)

存在的问题：

• self.images 的类型一定是 list 吗？其实可以是 numpy 数组
• self.images 的每个元素，和其他元素，一定是独立的吗？ 可以是同一个内存上连续的分布
• self.images 的每个元素，内存可以和读取文件得到的 buffer 复用吗？可以！

class DataSet:def __init__(self):self.images = Noneself.labels = Noneself.buf = Nonedef load_images(self, filename):with open(filename, 'rb') as fin:self.buf = bytearray(fin.read())magic = np.frombuffer(self.buf, np.uint8, count=4)num_dims = magic[3]dims = np.frombuffer(self.buf, dtype=np.dtype('>u4'), count=num_dims, offset=4)num_images, rows, cols = dimsself.images = np.frombuffer(self.buf, dtype=np.uint8, offset=4+4*num_dims).reshape(num_images, rows, cols)...train_set = DataSet()
train_set.load_images('data/train-images.idx3-ubyte')
print("Images and buffer share memory:", np.shares_memory(train_set.images, train_set.buf))

解释：self.buf 的类型，如果直接用 fin.read() 则得到 bytes 类型，是不可变的；转为 bytearray 类型后，是可变的，就可以保持和 self.images( ) 共享。

遗憾的是， self.buf = bytearray(fin.read()) 这句本身就发生了内存拷贝。

改进 - 避免内存拷贝

with open(filename, 'rb') as fin:  self.buf = bytearray(fin.read())  # 当前实现，存在两次内存分配  

改为

with open(filename, 'rb') as fin:  self.buf = fin.read()  # 读取为 bytes  self.buf = memoryview(self.buf)  # 直接使用 memoryview  

就可以避免 bytes 对象的中间拷贝过程。

完整代码

import numpy as np
import cv2class DataSet:def __init__(self):self.images = Noneself.labels = Noneself.buf = Nonedef load_images(self, filename):with open(filename, 'rb') as fin:#self.buf = bytearray(fin.read())self.buf = fin.read()self.buf = memoryview(self.buf)magic = np.frombuffer(self.buf, np.uint8, count=4)num_dims = magic[3]dims = np.frombuffer(self.buf, dtype=np.dtype('>u4'), count=num_dims, offset=4)num_images, rows, cols = dimsself.images = np.frombuffer(self.buf, dtype=np.uint8, offset=4+4*num_dims).reshape(num_images, rows, cols)def load_labels(self, filename):with open(filename, 'rb') as fin:buf = fin.read()magic = np.frombuffer(buf, np.uint8, count=4)num_dims = magic[3]dims = np.frombuffer(buf, dtype=np.dtype('>u4'), count=num_dims, offset=4)num_labels = dims[0]assert num_labels == len(self.images)self.labels = np.frombuffer(buf, dtype=np.uint8, offset=4+4*num_dims)def show_image(self, index):cv2.imshow('image', self.images[index])print('label:', self.labels[index])cv2.waitKey(0)cv2.destroyAllWindows()def main():train_set = DataSet()train_set.load_images('data/train-images.idx3-ubyte')train_set.load_labels('data/train-labels.idx1-ubyte')# train_set.show_image(0)# train_set.show_image(2)# train_set.show_image(5)print("Images and buffer share memory:", np.shares_memory(train_set.images, train_set.buf))if __name__ == '__main__':main()

总结

在前一篇，我们解析了MNIST数据集的IDX格式并用C++做了文件读取的实现，在本篇则切换到 Python 语言，在降低70%代码量的情况下实现了相同功能，并且避免了不必要的内存拷贝。这份工程之美，建立在对 IDX 格式有所了解的前提之下，对于 Python 的熟悉也是必不可少的，对于C++的经验也促使了复用内存这一条件的达成。