【知识】对比Share mem/Pin mem/GPU mem之间的传输速度

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn]

如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

目录

参考代码

运行结果

---

参考代码

import torch
import time
import matplotlib.pyplot as plt# 初始化设备和张量
device = torch.device('cuda')
data_sizes = [100, 1000, 5000, 10000, 50000, 100000, 300000, 500000]  # 不同数据量
results = {'Shared to Pinned': [], 'Pinned to Shared': [],'GPU to Pinned': [], 'Pinned to GPU': [],'GPU to Shared': [], 'Shared to GPU': []}# 测试不同数据量
for size in data_sizes:shared_tensor = torch.randn((size, 1000), dtype=torch.float32, device='cpu').share_memory_()pinned_tensor = torch.randn((size, 1000), dtype=torch.float32, device='cpu').pin_memory()gpu_tensor = torch.randn((size, 1000), dtype=torch.float32, device=device)# Shared Memory => Pinned Memorystart_time = time.time()pinned_tensor.copy_(shared_tensor, non_blocking=True)end_time = time.time()results['Shared to Pinned'].append(end_time - start_time)torch.cuda.synchronize()# Pinned Memory => Shared Memorystart_time = time.time()shared_tensor.copy_(pinned_tensor, non_blocking=True)end_time = time.time()results['Pinned to Shared'].append(end_time - start_time)torch.cuda.synchronize()# GPU Memory => Pinned Memorystart_time = time.time()pinned_tensor.copy_(gpu_tensor, non_blocking=True)end_time = time.time()results['GPU to Pinned'].append(end_time - start_time)torch.cuda.synchronize()# Pinned Memory => GPU Memorystart_time = time.time()gpu_tensor.copy_(pinned_tensor, non_blocking=True)end_time = time.time()results['Pinned to GPU'].append(end_time - start_time)torch.cuda.synchronize()# GPU Memory => Shared Memorystart_time = time.time()shared_tensor.copy_(gpu_tensor, non_blocking=True)end_time = time.time()results['GPU to Shared'].append(end_time - start_time)torch.cuda.synchronize()# Shared Memory => GPU Memorystart_time = time.time()gpu_tensor.copy_(shared_tensor, non_blocking=True)end_time = time.time()results['Shared to GPU'].append(end_time - start_time)torch.cuda.synchronize()# 绘制图表
plt.figure(figsize=(10, 6))
for key, values in results.items():plt.plot(data_sizes, values, marker='o', label=key)plt.xlabel('Data Size (rows)')
plt.ylabel('Time (seconds)')
plt.title('Memory Copy Time for Different Data Sizes')
plt.legend()
plt.grid(True)
plt.show()

运行结果

        所以，share to gpu是最慢的，而对于pin和gpu之间的互传非常快(异步传输)。以后如何选，心里也大概有个数了。