pytorch 使用gpu

原创于 2020-03-12 16:27:23 发布 · 438 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

pytorch 专栏收录该内容

7 篇文章

订阅专栏

本文详细介绍如何在PyTorch中检查和设置GPU设备，包括查看GPU是否可用、获取GPU数量、名称及当前设备索引，并提供代码示例说明如何使用device属性高效地将模型和数据置于GPU上进行计算。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1. pytorch查看设备信息

# 查看gpu设备是否可用
torch.cuda.is_available() 

# 返回gpu数量；
torch.cuda.device_count()

# 返回gpu名字，设备索引默认从0开始；
torch.cuda.get_device_name(0)

# 返回当前设备索引；
torch.cuda.current_device()

2. 设置device

一般我们设置第一个gpu（主要是因为我只有一块gpu。。。）

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

然后我们就可以使用to方法快捷的将数据放在gpu中计算：

model = CNN().to(device) # model放入gpu
criteon = nn.CrossEntropyLoss().to(device) # 损失函数
data, target = data.to(device), target.cuda() # 训练的数据

这里建议使用 to(device)，而不是 cuda()，毕竟 device 全局设置可以方便管理，且对于没有gpu的设备，device可以设置为cpu。