non_blocking=True 参数，imgs = imgs.to(device, non_blocking=True)

non_blocking=True 参数的作用与数据从主机（CPU）到设备（GPU）的异步传输有关：

---

1. 什么是非阻塞（Non-blocking）传输？

• 当数据从主机内存（CPU）传输到设备内存（GPU）时，默认情况下，这种传输是阻塞的。也就是说，当前操作会等待数据传输完成后，才会继续执行其他操作。
• 设置 non_blocking=True 后，如果条件满足，数据传输会变成异步操作。这样主机可以在传输数据的同时执行其他任务，从而提高程序的整体效率。

---

2. 条件：什么情况下 non_blocking=True 生效？

• 数据必须存储在主机的“固定内存”（Pinned Memory）中：

  • Pinned memory 是一种特殊类型的主机内存，支持更快的 GPU 数据传输。
  • 如果数据不在固定内存中，即使 non_blocking=True，传输仍然会是同步的。
  • 可以通过 torch.utils.data.DataLoader 设置 pin_memory=True，使加载的数据位于固定内存中。

    DataLoader(dataset, batch_size=32, pin_memory=True)
    

• 目标设备必须是 GPU：

  • non_blocking 参数只在主机（CPU）到设备（GPU）的传输中有效。
  • 如果目标设备是 CPU，non_blocking=True 没有任何意义。

---

3. 使用 non_blocking=True 的场景

• 深度学习模型训练中：
  • 训练时，数据加载与 GPU 计算可以重叠进行。例如，在 GPU 处理当前 batch 的同时，下一 batch 的数据可以异步传输到 GPU。
  • 通过异步数据传输，可以避免数据传输成为训练的性能瓶颈。

---

4. 示例

import torch
from torch.utils.data import DataLoader, TensorDataset

# 创建数据集并设置 pin_memory=True
dataset = TensorDataset(torch.randn(1000, 3, 224, 224))  # 随机数据
dataloader = DataLoader(dataset, batch_size=32, pin_memory=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

for imgs in dataloader:
    # 非阻塞传输到 GPU
    imgs = imgs[0].to(device, non_blocking=True)
    # 模型推理或训练
    outputs = imgs.mean(dim=(1, 2, 3))  # 示例操作

---

5. 性能优化的注意事项

• 固定内存（Pinned Memory）：

  • 固定内存是 non_blocking=True 生效的前提，可以通过 DataLoader(pin_memory=True) 启用。
  • 固定内存会占用更多的主机资源，因此需要根据实际需求权衡是否启用。

• 避免无意义的设置：

  • 如果数据已经在 GPU 上或者传输到 CPU，non_blocking=True 会被忽略，不会有任何性能提升。

---

6. 总结

non_blocking=True 的主要作用是利用异步数据传输，提高 GPU 和 CPU 的计算并行效率。在大规模训练任务中，正确配置可以显著提升数据加载与传输的效率，但需要确保数据位于固定内存且目标设备是 GPU。