Python 自动寻找可用的GPU

Oysterman

已于 2023-04-03 21:57:50 修改

阅读量988

点赞数

CC 4.0 BY-SA版权

文章标签： python pytorch 深度学习

于 2023-04-02 16:49:05 首次发布

本文链接：https://blog.csdn.net/zhziqiang/article/details/129912886

使用包含多个GPU的主机时，每次可用的GPU不是固定的，需要手动指定使用哪个GPU。但是，其实可以，自动寻找可用的GPU。

实现1：基于PyTorch

# 参照 lightning (原 pytorch-lightning)
import torch

"""
寻找空闲的 GPU
"""
def find_availabel_gpu(nb=1):
    for i in range(torch.cuda.device_count(), -1, -1):
        gpu_mem_size = torch.cuda.get_device_properties(0).total_memory  # 字节
        gpu_mem_size = int(gpu_mem_size * 0.95)
        # Try to allocate on device:
        device = torch