查看设备GPU数目

最新推荐文章于 2024-09-02 23:02:16 发布

Matt今年18岁

最新推荐文章于 2024-09-02 23:02:16 发布

阅读量1.8k

点赞数

分类专栏：新手跑深度学习代码文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/weixin_44698281/article/details/121183473

版权

PyTorch GPU 张量设备分配 CUDA

关键词由CSDN通过智能技术生成

新手跑深度学习代码专栏收录该内容

7 篇文章 1 订阅

订阅专栏

import torch

def try_gpu(i=0):
    """如果存在返回gpu(i), 否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():
    """返回该设备上的所有GPU数"""
    devices = [
        torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

a = try_gpu()
b = try_gpu(1)
c = try_all_gpus()
print(a, b, c)