Pytorch调用GPU的方法

最新推荐文章于 2024-06-17 14:17:20 发布

最幸伏的人

最新推荐文章于 2024-06-17 14:17:20 发布

阅读量3.4k

点赞数 2

分类专栏：人工智能文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/qq_43631827/article/details/131629262

版权

人工智能专栏收录该内容

9 篇文章 1 订阅

订阅专栏

Pytorch调用GPU有两种方法，一种是torch.cuda()，一种是torch.to()。

torch.cuda()

通常会在配置文件中写入调用的GPU，默认不填参数为0
gpu = [0, 1, 2]
# 默认调用调用0号GPU
network = network.cuda()
# 指定调用0号GPU
network = network.cuda(gpu[0])
# 指定调用1号GPU
network = network.cuda(gpu[1])

torch.to()

GPU = ['cuda', 'cuda:0', 'cuda:1', 'cuda:2']
# 这个写法是最常见的写法，默认调用0号GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 指定调用0号GPU
device = torch.device(GPU[1] if torch.cuda.is_available() else "cpu")
# 指定调用1号GPU
device = torch.device(GPU[2] if torch.cuda.is_available() else "cpu")

network = network.to(device)