Nothing can help us endure dark times better than our faith.
事先了解下「物理内存和虚拟内存」,有助于阅读和理解本文。
在 CUDA 编程中,内存拷贝是非常费时的一个动作。CPU 和 GPU 之间的总线是 PCI-Express,是双向传输的,之间通过 DMA(Direct Memory Access,直接内存访问)机制进行数据拷贝。
1.CUDA 固定内存(Pinned Memory)
对于 CUDA 架构而言,主机端的内存可分为两种:
- 1)可分页内存(Pageable Memory)。
- 2)页锁定内存(Page-locked Memory),或称固定内存(Pinned Memory)。
对于页锁定内存,操作系统不会对其进行分页和交换操作,一定是存储在物理内存,不会存储在虚拟内存,因此,GPU 可直接通过 DMA 机制,在主机和 GPU 之间快速复制数据。而对于可分页内存,CUDA 驱动会首先将 Pageable Memory 的数据拷贝到临时的 Pinned Memory,然后再通过 DMA 机制进行数据传输。如下图所示:
可见,从主机的页锁定内存拷贝数据到 GPU 要更快的多。
但是,页锁定内存也是更消耗物理内存的。当计算机的物理内存充足时,可以设置为页锁定内存。当计算机的物理内存不足或内存交换过多时,不建议使用页锁定内存,否则会严重影响应用程序的性能。
C++ 中,通过操作系统接口 malloc() 在主机上分配可分页内存,通过 CUDA 函数 cudaHostAlloc() 在主机上分配页锁定内存。
2.torch.utils.data.DataLoader() 中的 pin_memory 参数
若创建 DataLoader 时设置 pin_memory=True,则返回放置在页锁定内存中的批数据,使得将内存的 Tensor 数据拷贝到 GPU 显存变得更快。
因为 pin_memory 与电脑硬件性能有关,由于不能确保每一个炼丹玩家都有高端设备,因此PyTorch 官网默认 pin_memory=False。