本文将介绍如何应用内存映射文件放慢 PyTorch 数据集的加载速度
在应用 Pytorch 训练神经网络时，最常见的与速度相干的瓶颈是数据加载的模块。如果咱们将数据通过网络传输，除了预取和缓存之外，没有任何其余的简略优化形式。
然而如果数据本地存储，咱们能够通过将整个数据集组合成一个文件，而后映射到内存中来优化读取操作，这样咱们每次文件读取数据时就不须要拜访磁盘，而是从内存中间接读取能够放慢运行速度。
什么是内存映射文件
内存映射文件（memory-mapped file）是将残缺或者局部文件加载到内存中，这样就能够通过内存地址相干的 load 或者 store 指令来操纵文件。为了反对这个性能，古代的操作系统会提供一个叫做 mmap 的零碎调用。这个零碎调用会接管一个虚拟内存地址（VA），长度（len），protection，一些标记位，一个关上文件的文件描述符，和偏移量（offset）。
因为虚拟内存代表的附加形象层，咱们能够映射比机器的物理内存容量大得多的文件。正在运行的过程所需的内存段（称为页）从内部存储中获取，并由虚拟内存管理器主动复制到主内存中。
应用内存映射文件能够进步 I / O 性能，因为通过零碎调用进行的一般读 / 写操作比在本地内存中进行更改要慢得多，对于操作系统来说，文件以一种“惰性”的形式加载，通常一次只加载一个页，因而即便对于较大的文件，理论 RAM 利用率也是最低的，然而应用内存映射文件能够改善这个流程。
什么是 PyTorch 数据集
Pytorch 提供了用于在训练模型时解决数据管道的两个次要模块:Dataset 和 DataLoader。
DataLoader 次要用作 Dataset 的加载，它提供了许多可配置选项，如批处理、采样、预读取、变换等，并形象了许多办法。
Dataset 是咱们进行数据集解决的理论局部，在这里咱们编写训练时读取数据的过程，包含将样本加载到内存和进行必要的转换。
对于 Dataset，必须实现:
__init_,
__len__和
__getitem__三个办法
实现自定义数据集
接下来，咱们将看到下面提到的三个办法的实现。
最重要的局部是在
__init__中，咱们将应用 numpy 库中的
np.memmap()函数来创立一个 ndarray 将内存缓冲区映射到本地的文件。
在数据集初始化时，将 ndarray 应用可迭代对象进行填充，代码如下：
class MMAPDataset(Dataset):
def __init__(
self,
input_iter: Iterable[np.ndarray],
labels_iter: Iterable[np.ndarray],
mmap_path: str = None,
size: int = None,
transform_fn: Callable[..., Any] = None,
) -> None:
super().__init__()
self.mmap_inputs: np.ndarray = None
self.mmap_labels: np.ndarray = None
self.transform_fn = transform_fn
if mmap_path is None:
mmap_path = os.path.abspath(os.getcwd())
self._mkdir(mmap_path)
self.mmap_input_path = os.path.join(mmap_path, DEFAULT_INPUT_FILE_NAME)
self.mmap_labels_path = os.path.join(mmap_path, DEFAULT_LABELS_FILE_NAME)
self.length = size
for idx, (input, label) in enumerate(zip(input_iter, labels_iter)):
if self.mmap_inputs is None:
self.mmap_inputs = self._init_mmap(self.mmap_input_path, input.dtype, (self.length, *input.shape)
)
self.mmap_labels = self._init_mmap(self.mmap_labels_path, label.dtype, (self.length, *label.shape)
)
self.mmap_inputs[idx][:] = input[:]
self.mmap_labels[idx][:] = label[:]
def __getitem__(self, idx: int) -> Tuple[Union[np.ndarray, torch.Tensor]]:
if self.transform_fn:
return self.transform_fn(self.mmap_inputs[idx]), torch.tensor(self.mmap_labels[idx])
return self.mmap_inputs[idx], self.mmap_labels[idx]
def __len__(self) -> int:
return self.length咱们在下面提供的代码中还应用了两个辅助函数。
def _mkdir(self, path: str) -> None:
if os.path.exists(path):
return
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
return
except:
raise ValueError("Failed to create the path (check the user write permissions)."
)
def _init_mmap(self, path: str, dtype: np.dtype, shape: Tuple[int], remove_existing: bool = False) -> np.ndarray:
open_mode = "r+"
if remove_existing:
open_mode = "w+"
return np.memmap(
path,
dtype=dtype,
mode=open_mode,
shape=shape,
)能够看到，下面咱们自定义数据集与个别状况的次要区别就是
_init_mmap中调用的
np.memmap(), 所以这里咱们对
np.memmap()做一个简略的解释：
Numpy 的 memmap 对象，它容许将大文件分成小段进行读写，而不是一次性将整个数组读入内存。memmap 也领有跟一般数组一样的办法，基本上只有是能用于 ndarray 的算法就也能用于 memmap。
应用函数
np.memmap并传入一个文件门路、数据类型、形态以及文件模式，即可创立一个新的 memmap 存储在磁盘上的二进制文件创建内存映射。
对于更多的介绍请参考 Numpy 的文档，这里就不做具体的解释了
基准测试
为了理论展现性能晋升，我将内存映射数据集实现与以经典形式读取文件的一般数据集实现进行了比拟。这里应用的数据集由 350 张 jpg 图像组成。
从上面的后果中，咱们能够看到咱们的数据集比一般数据集快 30 倍以上：
总结
本文中介绍的办法在减速 Pytorch 的数据读取是十分无效的，尤其是应用大文件时，然而这个办法须要很大的内存，在做离线训练时是没有问题的，因为咱们可能齐全的管制咱们的数据，然而如果想在生产中利用还须要思考应用，因为在生产中有些数据咱们是无法控制的。
https://avoid.overfit.cn/post/33d9496e1f8440d69a220fe6b9ab700c
作者：Tudor Surdoiu