基于nvidia-ml-py实现深度学习过程中GPU的自动选择

1 基本介绍

1.1 契机

如果你采用的是英伟达的显卡来运行深度学习的代码，那么就需要选择指定的GPU来运行代码，一般情况是通过nvidia-smi来查看GPU的使用情况，然后选择空闲内存最大的那一张显卡进行代码的运行，同时代码可能还会在本机和服务器上来回调试，每次都需要自己来设置不同的device，可能会引起不必要的麻烦，而且每个主机的显卡数量是有限的，不可能不断的查看显卡的使用情况，如何在调用自己的代码进行运行。

因此就想着有没有办法来自动化的选择最合适的device，这样就不需要每次进行设置了，同时还可以不断的监听显卡的使用情况，在合适的时间点运行代码。

1.2 想法

在网上找寻相关的资料，找到了pynvml这个库，但是这个库并不是官方出品的，同时更新的频率也有点低，因此又在英伟达管理库中找到了nvidia-ml-py，该python库的最新更新时间是2024年4月，因此选择后者来解决上面的问题。

要实现寻找空闲内存最大的显卡，首先就需要知道每章显卡的使用情况，之后将每张显卡的空闲大小按顺序排序即可。选择空闲内存最大的显卡，将其空闲内存和预估的内存进行比较，如果空闲内存超过了预估内存则运行，否则则等待。

2 实现

2.1 代码实现

1. 环境安装

   pip install nvidia-ml-py

2. GPU内存监听

   import time
   import pynvml
   
   from managers.logger_manager import LoggerManager
   
   
   class GPUMonitor():
   
       def __init__(self, config: dict, logger: LoggerManager) -> None:
           self.config = config
           self.logger = logger
           try:
               pynvml.nvmlInit()
               self.driver_version = pynvml.nvmlSystemGetDriverVersion()
               self.cuda_version = pynvml.nvmlSystemGetCudaDriverVersion()
               self.num_device = pynvml.nvmlDeviceGetCount()
               self.gpu_info = {
                   "driver_version": "None",
                   "cuda_version": "None",
                   "num_device": "None",
                   "device": [],
               }
               self.run()
           except:
               self.config["device"] = -1
               self.logger.info("未检测到GPU设备")
               self.logger.send_message("未检测到GPU设备，使用CPU", message_type=1, message_content_type=1)
   
       def get_gpu_info(self):
           self.gpu_info['driver_version'] = self.driver_version
           self.gpu_info['cuda_version'] = self.cuda_version
           self.gpu_info['num_device'] = self.num_device
           for idx in range(self.num_device):
               handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
               device_name = pynvml.nvmlDeviceGetName(handle)
               mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
               total_mem = int(int(mem_info.total) / 1024 / 1024)
               used_mem = int(int(mem_info.used) / 1024 / 1024)
               free_mem = int(int(mem_info.free) / 1024 / 1024)
               util = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
               temp = pynvml.nvmlDeviceGetTemperature(handle, 0)
               self.gpu_info["device"].append({
                   "device_name": device_name,
                   "idx": idx,
                   "temp": temp,
                   "used_mem": used_mem,
                   "free_mem": free_mem,
                   "total_mem": total_mem,
                   "gpu_util": util,
               })
               self.gpu_info["device"].sort(key=lambda x: x["free_mem"], reverse=True)
   
       def run(self):
           cycle_num = self.config.get("cycle_num", 30)
           while cycle_num > 0:
               self.get_gpu_info()
               if self.gpu_info["num_device"] == 0:
                   self.config["device"] = -1
                   self.logger.info("未检测到GPU设备")
                   self.logger.send_message("未检测到GPU设备，使用CPU", message_type=1, message_content_type=1)
                   break
               need_free_mem = self.config.get("need_free_mem", 1000)
               if self.gpu_info["device"][0]["free_mem"] > need_free_mem:
                   self.config["device"] = self.gpu_info["device"][0]["idx"]
                   self.logger.info(f"选择GPU: {self.gpu_info['device'][0]['idx']}, 可用内存: {self.gpu_info['device'][0]['free_mem']}MB")
                   self.logger.send_message(f"选择GPU: {self.gpu_info['device'][0]['idx']}, 可用内存: {self.gpu_info['device'][0]['free_mem']}MB", message_type=2, message_content_type=1)
                   break
               else:
                   self.logger.info("GPU资源不足，等待10分钟后重新检查")
                   self.logger.send_message("GPU资源不足，等待10分钟后重新检查", message_type=1, message_content_type=1)
                   time.sleep(600)
                   cycle_num -= 1

   • config传入的是代码中使用的config，直接就config进行修改即可
   • logger传入的是日志模块，将相关的日志输出

3. 在config中使用

   from config.data_config import DATA_CONFIG
   from config.global_config import GOLBAL_CONFIG
   from managers.logger_manager import LoggerManager
   from utils.gpu_utils import GPUMonitor
   
   
   class ConfigManager():
   
       def __init__(self, train_config):
           self.train_config = train_config
           self.data_config = {}
           self.default_train_config = {
               "debug_mode": True,
               "epoch": 5,
               "batch_size": 1024,
               "lr": 0.0001,
               "device": -1,
               "metric_func": {
                   "train": [
                       {
                           "eval_func": "auc"
                       },
                       {
                           "eval_func": "log_loss"
                       },
                   ],
                   "eval": [
                       {
                           "eval_func": "ndcg",
                           "k": [10, 20, 50]
                       },
                       {
                           "eval_func": "gauc"
                       },
                   ],
               },
           }
           self.global_config = {}
           self.required_params = ["data", "model_name", "trainer"]
   
       def _get_data_config(self):
           self.data_config = DATA_CONFIG.get(self.train_config.get('data'))
           if self.data_config is None:
               raise ValueError(f"不支持的数据集: {self.train_config['data']}")
   
       def _get_global_config(self):
           self.global_config = GOLBAL_CONFIG
   
       def _check_params(self):
           missing_params = [param for param in self.required_params if param not in self.train_config.keys()]
           if len(missing_params) > 0:
               raise ValueError(f"缺少必要参数: {missing_params}")
           for param in self.default_train_config.keys():
               if param not in self.train_config.keys():
                   self.train_config[param] = self.default_train_config[param]
                   self.logger.warning(f"缺少可选参数 {param}，设置默认值: {self.default_train_config[param]}")
   
       def get_config(self):
           self._get_data_config()
           self._get_global_config()
           self.logger = LoggerManager(config=self.global_config)
           # 在此处进行使用
           GPUMonitor(config=self.train_config, logger=self.logger)
           self._check_params()
           all_config = {**self.global_config, **self.data_config, **self.train_config}
           self.logger.send_message({**self.data_config, **self.train_config}, message_type=0, message_content_type=0)
           return all_config

4. 因为dict是可变对象，因此对其进行增删改查的情况下，所有地方都会一起变化，因此没有重新进行赋值

2.2 效果展示

def set_device(device: int) -> torch.device:
    """设置设备

    Args:
        device (int): 设备号

    Returns:
        torch.device: 设备
    """
    if device == -1:
        return torch.device("cpu")
    elif device == -2:
        return torch.device("mps")
    else:
        return torch.device(f"cuda:{device}")

本人的运行程序中只需要传入对应的数字即可选择对应的设备进行训练，因此在配置中只需要传入数字即可。

2.2.1 存在显卡

1. 使用nvidia-smi查看当前显卡运行情况

   

2. 配置train_config.py文件，使得device为1和3外的卡

   

3. 运行项目代码，打断点进行测试

   

2.2.2 不存在显卡

1. 本人的开发机器是MAC的m1芯片，因此不存在显卡，正常会设置为device设置为-1或者-2

2. 配置train_config.py文件，使得device为非负数

   

3. 运行项目代码，打断点进行测试，可以看见device变为了-1

   

2.2.3 显卡占满

1. 在程序运行之前可以先运行一个batch_size来大概确定程序所需要的内存大小。
2. 我个人会在一个batch_size所需要内存的基础上乘以1.5作为配置传入need_free_mem中。
3. 程序中默认设置的need_free_mem是1000MB，默认设置的循环轮次是30次，默认设置的等待时长是10分钟，可以按照自己的需要进行调整。

3 未来

1. 同时并行运行两个模型也可以为每个模型选择合适的显卡，目前的想法是直接采用多线程，但是每个模型运行的中间有一个间隔，来确定上一个模型使用的内存大小。
2. 配置多张显卡来并行运行单个模型，目前的想法是在config中出传入并行运行模型的显卡数量，然后使用nn.DataParallel来运行。