pytorch优化memo

1 pin-memory 锁页内存
CPU内存分配默认是可分页的,GPU无法直接从可分页CPU内存访问数据,因此需要先创建一个临时的缓冲区(pinned memory),把数据从可分页内存拷贝pinned内存上,然后GPU才能从pinned内存上读取数据,
通过设置pin_memory=True,从一开始就把一部分内存给锁住(上图(右)),从而减少了Host内部的开销,避免了CPU内存拷贝时间。
1 | train_sampler = None |
2 non_blocking 异步
从Host端(主机侧)来看,device 侧的操作可以分为:异步 和 同步 。
- 异步操作:在调用接口或者kernel launch后,立刻将控制权交还给CPU,不阻塞主机继续执行
- 核函数(kernel function)
- cuda的非阻塞接口,比如 cudaMemcpyAsync、cudaMemsetAsync、cudaEventQuery 等
- 同步操作:阻塞主机的执行,直至当前操作执行完毕。
- 内存操作接口:cudaMemcpy 、cudaMemset、cudaSetDevice 等
- 设备相关接口:cudaSetDevice、cudaGetDeviceProperties 等
1 | data = data.to('cuda:0', non_blocking=True) |
non_blocking=True用于指定是否启用异步数据传输。对于非默认stream,non-default stream上的所有操作相对于 host code 都是 non-blocking 的,即它们不会阻塞Host代码。意思是,等于可能是多流操作,可能是cudaMemcpyAsync(?)。