pytorch创建data.DataLoader时,参数pin_memory的理解

2,282 阅读2分钟

锁页内存理解(pinned memory or page locked memory):

主机端存在虚拟内存,主机内存不足是会将内存数据交换到虚拟内存中,虚拟内存就是主机中的磁盘空间,需要该页时再重新从磁盘加载回来。这样做可以使用比实际内存更大的内存空间。

锁页内存允许GPU上的MDA控制器在使用主机内存时不用CPU参与。GPU上的显存都是锁页的,因为GPU上的内存时不支持交换到磁盘的。锁页内存就是分配主机内存时锁定该页,让其不与磁盘交换

CUDA中锁页内存的使用可以使用CUDA驱动API( driver API’s)cuMemAllocHost()或者使用CUDA的运行时API(runtime API)中的cudaMallocHost()。除此之外还可以直接用主机上Malloc()分配的空间,然后将其注册为锁页内存(使用cudaHostRegister()函数完成注册)。

使用锁页内存的好处有以下几点:

  • 设备内存与锁页内存之间的数据传输可以与内核执行并行处理。
  • 锁页内存可以映射到设备内存,减少设备与主机的数据传输。
  • 在前端总线的主机系统锁页内存与设备内存之间的数据交换会比较快;并且可以是write-combining的,此时带宽会跟大。

pin_memory就是锁页内存,创建DataLoader时,设置pin_memory=True,则意味着生成的Tensor数据最开始是属于内存中的锁页内存,这样将内存的Tensor转义到GPU的显存就会更快一些。

主机中的内存,有两种存在方式,一是锁页,二是不锁页锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换(注:虚拟内存就是硬盘),而不锁页内存在主机内存不足时,数据会存放在虚拟内存中

而显卡中的显存全部是锁页内存!

当计算机的内存充足的时候,可以设置pin_memory=True。当系统卡住,或者交换内存使用过多的时候,设置pin_memory=False。因为pin_memory与电脑硬件性能有关,pytorch开发者不能确保每一个炼丹玩家都有高端设备,因此pin_memory默认为False