理解CPU与GPU频繁数据传输

基础理解

在学习深度学习神经网络过程中，有时候会遇到一些描述“尽量避免CPU与GPU频繁数据传输”。那这句话应该如何理解呢？我们知道CPU可以访问内存，而GPU也有自己的显存。要完成功能一般都是CPU从硬盘或者其他数据源读取数据到内存中，然后将内存中的传输到GPU的显存中，GPU从显存中获取数据并进行计算，并最终将计算的结果返回给CPU的内存中。

整体的计算就像上面描述，但是不可忽略的是：

1. 从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()
        
        # 只在需要显示或保存时才传回CPU
        if 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 上