基础理解
在学习深度学习神经网络过程中,有时候会遇到一些描述“尽量避免CPU与GPU频繁数据传输”。那这句话应该如何理解呢?我们知道CPU可以访问内存,而GPU也有自己的显存。要完成功能一般都是CPU从硬盘或者其他数据源读取数据到内存中,然后将内存中的传输到GPU的显存中,GPU从显存中获取数据并进行计算,并最终将计算的结果返回给CPU的内存中。
整体的计算就像上面描述,但是不可忽略的是:
- 从CPU内存到GPU显存之间的数据传输是开销的,也是延迟的。
如果频繁的进行传输,就会加大这个开销,举个形象的例子:现在CPU内存和GPU显存相当于两个盒子,假设CPU的盒子中装满了苹果,而GPU的盒子没有苹果,现在需要将CPU盒子内的苹果转移到GPU的盒子中,如果从CPU所在的盒子每次拿一个苹果放到GPU所在的盒子,那么有多少个苹果就需要重复多少次,每次这个过程都需要消耗时间;那如果每次拿10个苹果到GPU所在的盒子呢?重复次数是否大大的小减少了,而且效率也得到了提升。
通过这个例子我们可以得到这样一个结论:CPU和GPU之间的数据传递每次传递尽量多的数据(淡然也不是无限制的多,但是不妨碍我们这样理解)。
错误的示范
# 示例1:循环中的频繁传输
for i in range(1000):# 每次循环都进行CPU->GPU传输data = torch.tensor([i]).cuda() # ❌ 不好result = model(data)# GPU->CPU传输result = result.cpu().numpy() # ❌ 不好# 示例2:逐个处理数据
for image in dataset:# 每张图片单独传输image = image.cuda() # ❌ 不好output = model(image)output = output.cpu() # ❌ 不好
为什么要避免频繁传输?
1.带宽限制
# PCIe 4.0 x16 理论带宽约为 32 GB/s
# 但实际传输速度要低得多# 假设传输一个 batch 的数据 (32, 3, 224, 224) float32
data_size = 32 * 3 * 224 * 224 * 4 # ~19MB# 如果每个样本单独传输:
# 需要32次传输,每次产生延迟
# 总延迟 = 传输次数 * (基础延迟 + 传输时间)
2.延迟开销
- 每次传输都有固定的启动开销
- 多次小数据传输比一次大数据传输效率低
正确的做法
1.批量处理
# 好的实践:一次性传输整个批次
batch_data = torch.stack([torch.tensor(x) for x in data])
batch_data = batch_data.cuda() # ✅ 一次性传输
outputs = model(batch_data)
results = outputs.cpu() # ✅ 一次性传回
2.数据加载器优化
# 使用 DataLoader 进行批处理
dataloader = DataLoader(dataset, batch_size=32,pin_memory=True, # ✅ 使用固定内存num_workers=4) # ✅ 多进程加载# 训练循环
for batch in dataloader:batch = batch.cuda() # ✅ 每个批次只传输一次outputs = model(batch)# 除非必要,尽量在GPU上完成所有计算
3.保持数据在GPU上
# 好的实践:模型和数据都保持在GPU上
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())for epoch in range(num_epochs):for batch in dataloader:batch = batch.cuda()# 所有运算都在GPU上完成outputs = model(batch)loss = criterion(outputs, targets.cuda())loss.backward()optimizer.step()# 只在需要显示或保存时才传回CPUif i % display_step == 0:print(f"Loss: {loss.item()}") # .item() 只传输单个标量
常见的错误陷阱
隐式传输
# 隐式CPU-GPU传输
model.cuda()
for batch in dataloader:# 这里的运算符会触发隐式传输output = model(batch + 1) # ❌ batch 在 CPU 上# 正确做法
model.cuda()
for batch in dataloader:batch = batch.cuda()output = model(batch + 1) # ✅ 所有运算都在 GPU 上