【知识】pytorch中的pinned memory和pageable memory
转载请注明出处:小锋学长生活大爆炸[xfxuezhagn.cn]
如果本文帮助到了你,欢迎[点赞、收藏、关注]哦~
目录
概念简介
pytorch用法
速度测试
反直觉情况
概念简介
默认情况下,主机 (CPU) 数据分配是可分页的。GPU 无法直接从可分页主机内存访问数据,因此当调用从可分页主机内存到设备内存的数据传输时,CUDA 驱动程序必须首先分配一个临时的页面锁定或“固定”主机数组,将主机数据复制到固定数组,然后将数据从固定阵列传输到设备内存。
如图所示,固定内存用作从设备到主机的传输暂存区域。通过直接在固定内存中分配主机阵列,我们可以避免在可分页主机阵列和固定主机阵列之间传输的成本。使用 cudaMallocHost() 或 cudaHostAlloc() 在 CUDA C/C++ 中分配固定主机内存,并使用 cudaFreeHost() 解除分配。固定内存分配可能会失败,因此应始终检查错误。
数据传输速率可能取决于主机系统的类型(主板、CPU 和芯片组)以及 GPU。通过运行BandwidthTest
会产生以下结果。可见,固定传输的速度是可分页传输的两倍多。(我的测试发现,基本上能跑满PCIe的带宽。)
Device: NVS 4200M
Transfer size (MB): 16Pageable transfersHost to Device bandwidth (GB/s): 2.308439Device to Host bandwidth (GB/s): 2.316220Pinned transfersHost to Device bandwidth (GB/s): 5.774224Device to Host bandwidth (GB/s): 5.958834
不过,不应过度分配固定内存。这样做会降低整体系统性能,因为它会减少操作系统和其他程序可用的物理内存量。多少是太多是很难提前判断出来的,因此与所有优化一样,测试你的应用程序及其运行的系统以获得最佳性能参数。
用法示例
由于pinned memory后,可以使用DMA传输而不占用CPU,因此通常需要搭配non_blocking使用。
# tensor.pin_memory() 就行
pinned_tensor = torch.randn(data_size, dtype=torch.float32).pin_memory()device = torch.device("cuda")
pinned_tensor.to(device, non_blocking=True)
速度测试
import torch
import time
import torch.multiprocessing as mp# 数据大小
data_size = 10**7 # 例如,10M数据def test_pinned_memory(rank, normal_tensor, pinned_tensor, device):# 测试普通内存到GPU传输时间start_time = time.perf_counter()normal_tensor_gpu = normal_tensor.to(device, non_blocking=True)torch.cuda.synchronize() # 等待数据传输完成normal_memory_time = time.perf_counter() - start_timeprint(f"[进程 {rank}] 普通内存到GPU传输时间: {normal_memory_time:.6f} 秒")# 测试固定内存到GPU传输时间start_time = time.perf_counter()pinned_tensor_gpu = pinned_tensor.to(device, non_blocking=True)torch.cuda.synchronize() # 等待数据传输完成pinned_memory_time = time.perf_counter() - start_timeprint(f"[进程 {rank}] 固定内存到GPU传输时间: {pinned_memory_time:.6f} 秒")# 比较结果speedup = normal_memory_time / pinned_memory_timeprint(f"[进程 {rank}] 固定内存的传输速度是普通内存的 {speedup:.2f} 倍")if __name__ == '__main__':# 分配普通内存中的张量normal_tensor = torch.randn(data_size, dtype=torch.float32)# 分配固定内存中的张量pinned_tensor = torch.randn(data_size, dtype=torch.float32).pin_memory()# 目标设备device = torch.device("cuda")# 使用mp.spawn启动多进程测试mp.spawn(test_pinned_memory, args=(normal_tensor, pinned_tensor, device), nprocs=2, join=True)
输出:
[进程 0] 普通内存到GPU传输时间: 1.054590 秒
[进程 0] 固定内存到GPU传输时间: 0.012945 秒
[进程 0] 固定内存的传输速度是普通内存的 81.47 倍
[进程 1] 普通内存到GPU传输时间: 1.169124 秒
[进程 1] 固定内存到GPU传输时间: 0.013019 秒
[进程 1] 固定内存的传输速度是普通内存的 89.80 倍
可以看到速度还是非常快的。
反直觉情况
我再瞎试的过程中发现,如果将pinned memory放在一个class中,那么多进程时候,pinned memory的移动很慢。暂不清楚为什么。
示例代码(反例,仅供观看,请勿使用):
import torch
import torch.multiprocessing as mp
class PinnedMemoryManager:def __init__(self, data_size):self.data_size = data_sizeself.normal_tensor = Noneself.pinned_tensor = Nonedef allocate_normal_memory(self):# 分配普通内存self.normal_tensor = torch.randn(self.data_size, dtype=torch.float32)def allocate_pinned_memory(self):# 分配固定内存self.pinned_tensor = torch.randn(self.data_size, dtype=torch.float32).pin_memory()def transfer_to_device(self, device, use_pinned_memory=False):# 选择使用普通内存或固定内存tensor = self.pinned_tensor if use_pinned_memory else self.normal_tensorif tensor is None:raise ValueError("Tensor not allocated. Call allocate_memory first.")# 数据传输start_time = torch.cuda.Event(enable_timing=True)end_time = torch.cuda.Event(enable_timing=True)start_time.record()tensor_gpu = tensor.to(device, non_blocking=True)end_time.record()# 同步并计算传输时间torch.cuda.synchronize()transfer_time = start_time.elapsed_time(end_time) / 1000.0 # 转换为秒return tensor_gpu, transfer_timedef free_memory(self):# 释放内存del self.normal_tensordel self.pinned_tensorself.normal_tensor = Noneself.pinned_tensor = Nonedef test_pinned_memory(rank, manager, device):# 测试普通内存到GPU传输时间normal_gpu, normal_memory_time = manager.transfer_to_device(device, use_pinned_memory=False)print(f"[进程 {rank}] 普通内存到GPU传输时间: {normal_memory_time:.6f} 秒")# 测试固定内存到GPU传输时间pinned_gpu, pinned_memory_time = manager.transfer_to_device(device, use_pinned_memory=True)print(f"[进程 {rank}] 固定内存到GPU传输时间: {pinned_memory_time:.6f} 秒")# 比较结果speedup = normal_memory_time / pinned_memory_timeprint(f"[进程 {rank}] 固定内存的传输速度是普通内存的 {speedup:.2f} 倍")if __name__ == '__main__':# 数据大小data_size = 10**7 # 例如,10M数据# 初始化固定内存管理器manager = PinnedMemoryManager(data_size)manager.allocate_normal_memory()manager.allocate_pinned_memory()# 目标设备device = torch.device("cuda")# 使用mp.spawn启动多进程测试mp.spawn(test_pinned_memory, args=(manager, device), nprocs=2, join=True)# 释放内存manager.free_memory()
输出:
[进程 1] 普通内存到GPU传输时间: 0.013695 秒
[进程 1] 固定内存到GPU传输时间: 0.013505 秒
[进程 1] 固定内存的传输速度是普通内存的 1.01 倍
[进程 0] 普通内存到GPU传输时间: 0.013752 秒
[进程 0] 固定内存到GPU传输时间: 0.013593 秒
[进程 0] 固定内存的传输速度是普通内存的 1.01 倍
可以看到基本上没啥改进。
暂不清楚原因,只能先无脑避免这种用法了。