pytorch优化memo
frc99 中本葱

1 pin-memory 锁页内存

pin-memory
CPU内存分配默认是可分页的,GPU无法直接从可分页CPU内存访问数据,因此需要先创建一个临时的缓冲区(pinned memory),把数据从可分页内存拷贝pinned内存上,然后GPU才能从pinned内存上读取数据,

通过设置pin_memory=True,从一开始就把一部分内存给锁住(上图(右)),从而减少了Host内部的开销,避免了CPU内存拷贝时间。

1
2
3
4
5
6
7
8
9
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)

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(?)。


摘自https://zhuanlan.zhihu.com/p/477870660

由 Hexo 驱动 & 主题 Keep