non_blocking=True 与 torch.cuda.synchronize()

需要注意的是: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

二、 torch.cuda.synchronize()

目的:保证执行完gpu的任务再执行下面的操作

这是因为:GPU和CPU之间是异步执行的,CPU向GPU下达指令以后会立刻执行之后的代码,所以我们需要使用`torch.cuda.synchronize() `显示地等待GPU完成运算。

有一个例子是这样的:

pytorch 正确的测试时间的代码 torch.cuda.synchronize()_torch.cuda.synchronize() time-CSDN博客

需要注意的是任何在CPU上执行的操作都会打断gpu的异步操作,会将其变成同步操作,例如print()

所以不必要太在乎 torch.cuda.synchronize(),这种东西主要是为了保证让程序变成同步执行

Pytorch中x.cuda(non_blocking=True)参数解释_x.cuda()-CSDN博客

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Pengsen Ma

太谢谢了

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值