torch.device和torch.cuda.set_device()
2. 如果使用torch.cuda.set_device('cuda:0')的方式,则改用。3. 即前者需要采用to()方法,而后者可以使用cuda()方法,而无需传递参数。1. 使用torch.device的方式设置显卡之后,需用使用。的方式将模型和数据放到GPU上。的方式将模型和数据加载到显卡中。
·
# 快速介绍
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")和torch.cuda.set_device("cuda:0")是常用的设置显卡的两种方式,二者的区别在于:
1. 使用torch.device的方式设置显卡之后,需用使用
import torch
device = torch.device("cuda:0")
model.to(device)
data = data.to(device)
的方式将模型和数据放到GPU上
2. 如果使用torch.cuda.set_device('cuda:0')的方式,则改用
import torch
torch.cuda.set_device("cuda:0")
model.cuda()
data = data.cuda()
的方式将模型和数据加载到显卡中
3. 即前者需要采用to()方法,而后者可以使用cuda()方法,而无需传递参数
欢迎来到FlagOS开发社区,这里是一个汇聚了AI开发者、数据科学家、机器学习爱好者以及业界专家的活力平台。我们致力于成为业内领先的Triton技术交流与应用分享的殿堂,为推动人工智能技术的普及与深化应用贡献力量。
更多推荐
所有评论(0)