non_blocking=True 与 torch.cuda.synchronize()

Pengsen Ma

已于 2024-06-22 09:46:49 修改

阅读量850

点赞数 9

分类专栏： python异常错误机器学习基础文章标签： python pytorch

于 2024-06-22 09:45:17 首次发布

本文链接：https://blog.csdn.net/weixin_43135178/article/details/139874111

版权

机器学习基础同时被 2 个专栏收录

195 篇文章

订阅专栏

python异常错误

109 篇文章

订阅专栏

需要注意的是：GPU和CPU之间是异步执行的，CPU向GPU下达指令以后会立刻执行之后的代码，CPU不会等待GPU执行完成

一、non_blocking=True

目的：压缩gpu的效果，避免CPU与GPU设备数据传输时间开销带来的计算效率低下

在 PyTorch 中， non_blocking 是一个布尔类型的参数，用于指定是否启用异步数据传输。当你将 Tensor 数据从 CPU 移动到 GPU 时，可以通过设置 non_blocking=True 来启用异步传输模式，从而提高计算效率。这是因为深度学习很多程序在执行的时候加载数据的时间非常久，常见的两种策略分别是：pin_memory 和 non_blocking

这里给出一个示例：

images = images.to(device, non_blocking=True)

# 标准化等操作
mean = torch.Tensor([0.978, 0.977, 0.979]).cuda().unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
std = torch.Tensor([0.131, 0.136, 0.129]).cuda().unsqueeze(0).unsqueeze(-1).unsqueeze(-1)

# 输入model
image = (images - mean) / std
y = model(image)

# 等待数据异步传输完成
torch.cuda.synchronize()

# logger更新等操作

这里是将image从CPU转到gpu的时候，同时使用non_blocking=True异步操作，可以让程序同时执行下面的标准化等操作，其中标准化等操作并不涉及对加载的image进行操作，所以异步执行没有问题，并不需要torch.cuda.synchronize()等待上述程序执行完成。

Q：这里会有一个疑惑？为什么可以在异步的情况下直接将images输入进model？

A：其实第一步的CPU转到gpu的速度是很快的，程序可以保证在标准化等操作的时候就执行完CPU转到gpu的传输，这里使用non_blocking=True 只是为了进一步压缩gpu的使用效率。

然后最后会执行一步 torch.cuda.synchronize()，这是因为下面是logger等操作，我们还是需要等model执行完再logger记录，否则会因为GPU和CPU之间是异步执行的，CPU不会等待GPU执行完成（model（image））执行完成就记录logger