【知识】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 倍

可以看到基本上没啥改进。

暂不清楚原因，只能先无脑避免这种用法了。