转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn] 如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~
目录
概念简介
pytorch用法
速度测试
反直觉情况
概念简介
默认情况下，主机 （CPU） 数据分配是可分页的。GPU 无法直接从可分页主机内存访问数据，因此当调用从可分页主机内存到设备内存的数据传输时，CUDA 驱动程序必须首先分配一个临时的页面锁定或“固定”主机数组，将主机数据复制到固定数组，然后将数据从固定阵列传输到设备内存。
如图所示，固定内存用作从设备到主机的传输暂存区域。通过直接在固定内存中分配主机阵列，我们可以避免在可分页主机阵列和固定主机阵列之间传输的成本。使用 cudaMallocHost（） 或 cudaHostAlloc（） 在 CUDA C/C 中分配固定主机内存，并使用 cudaFreeHost（） 解除分配。固定内存分配可能会失败，因此应始终检查错误。
数据传输速率可能取决于主机系统的类型（主板、CPU 和芯片组）以及 GPU。通过运行BandwidthTest会产生以下结果。可见，固定传输的速度是可分页传输的两倍多。（我的测试发现，基本上能跑满PCIe的带宽。）
Device: NVS 4200M
Transfer size (MB): 16
Pageable transfers
Host to Device bandwidth (GB/s): 2.308439
Device to Host bandwidth (GB/s): 2.316220
Pinned transfers
Host to Device bandwidth (GB/s): 5.774224
Device to Host bandwidth (GB/s): 5.958834
不过，不应过度分配固定内存。这样做会降低整体系统性能，因为它会减少操作系统和其他程序可用的物理内存量。多少是太多是很难提前判断出来的，因此与所有优化一样，测试你的应用程序及其运行的系统以获得最佳性能参数。
用法示例
由于pinned memory后，可以使用DMA传输而不占用CPU，因此通常需要搭配non_blocking使用。
代码语言：javascript复制# 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)
速度测试
代码语言：javascript复制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_time
print(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_time
print(f"[进程 {rank}] 固定内存到GPU传输时间: {pinned_memory_time:.6f} 秒")
# 比较结果
speedup = normal_memory_time / pinned_memory_time
print(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的移动很慢。暂不清楚为什么。
示例代码（反例，仅供观看，请勿使用）：
代码语言：javascript复制import torch
import torch.multiprocessing as mp
class PinnedMemoryManager:
def __init__(self, data_size):
self.data_size = data_size
self.normal_tensor = None
self.pinned_tensor = None
def 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_tensor
if 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_time
def free_memory(self):
# 释放内存
del self.normal_tensor
del self.pinned_tensor
self.normal_tensor = None
self.pinned_tensor = None
def 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_time
print(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 倍
可以看到基本上没啥改进。
暂不清楚原因，只能先无脑避免这种用法了。
