pin_memory 和 non_blocking的作用分别是什么?网上看了很多解释,只是稀里糊涂的有个感觉,就是用了这玩意速度能变快,但是不知所以然,这篇文章希望能帮助你解惑,也给自己做个笔记,以备日后查阅。

train_sampler = None
train_loader = torch.utils.data.DataLoader(
	train_dataset,
		...,
		pin_memory=True
)

for data, labels in train_loader:
	data = data.to('cuda:0', non_blocking=True)

1. pin_memory

1.1 什么是锁页内存(Pinned Memory/PageLocked Memory)?什么是"Pinned"?

参考Cuda锁页内存和零复制

通常我们的主机处理器是支持虚拟内存系统的,即使用硬盘空间来代替内存。大多数系统中虚拟内存空间被划分成许多页,它们是寻址的单元,页的大小至少是4096个字节。虚拟寻址能使一个连续的虚拟地址空间映射到物理内存并不连续的一些页。

如果某页的物理内存被标记为换出状态,它就可以被更换到磁盘上,也就是说被踢出内存了。如果下次需要该页了,则重新加载到内存里。显然如果这一页切换的非常频繁,那么会浪费不少时间。

锁页(pinned page)是操作系统常用的操作,就是为了使硬件外设直接访问CPU内存,从而避免过多的复制操作。被锁定的页面会被操作系统标记为不可被换出的,所以设备驱动程序给这些外设编程时,可以使用页面的物理地址直接访问内存,CPU也可以访问上述锁页内存,但是此内存是不能移动或换页到磁盘上的。另外,在GPU上分配的内存默认都是锁页内存,这只是因为GPU不支持将内存交换到磁盘上。

1.2 什么时候设置pin_memory=True?

参考 How to Optimize Data Transfers in CUDA C/C++

总结一下上一小节的内容就是:

内存可以分为 没锁的(pageable,可分页的) 和 锁了的(pinned)。

Host(例如CPU)的数据分配默认是pageable(可分页的),但是GPU是没法直接读取pageable内存里的数据的,所以需要先创建一个临时的缓冲区(pinned memory),把数据从pageable内存拷贝pinned内存上,然后GPU才能从pinned内存上读取数据,如下图(左)所示。

image.png

但是CPU将数据从pageable 内存拷贝到 临时的 pinned 内存是有时间开销的,而且这个pinned 内存 还只是临时的,所以用完之后会被销毁。所以为了进一步提高效率,我们需要设置pin_memory=True,作用就是从一开始就把一部分内存给锁住(上图(右)),这样一来就减少了Host内部的开销,避免了CPU内存拷贝时间。

按照官方的建议[1]是你默认设置为True就对了。

image.png

2. non_blocking

2.1 CUDA Default Streams