空港防疫算法交互
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

31 line
1.9KB

  1. import GPUtil
  2. # order- 确定返回可用 GPU 设备 ID 的顺序。order应指定为以下字符串之一:
  3. # 'first'- 按升序排列可用的 GPU 设备 ID(默认)
  4. # 'last'- 按 id 降序排列可用的 GPU 设备 id
  5. # 'random'- 随机订购可用的 GPU 设备 ID
  6. # 'load'- 按负载递增排序可用的 GPU 设备 ID
  7. # 'memory'- 通过升序内存使用来排序可用的 GPU 设备 ID
  8. # limit- 将返回的 GPU 设备 ID 数量限制为指定数量。必须是正整数。(默认 = 1)
  9. # maxLoad- 被认为可用的 GPU 的最大当前相对负载。负载大于 的 GPUmaxLoad不会返回。(默认 = 0.5)
  10. # maxMemory- 被视为可用的 GPU 的最大当前相对内存使用量。maxMemory不返回当前内存使用量大于的 GPU 。(默认 = 0.5)
  11. # includeNan- 真/假标志,指示是否包括负载或内存使用为 NaN 的 GPU(指示无法检索使用情况)。(默认 = 假)
  12. # excludeID- ID 列表,应从可用 GPU 列表中排除。见GPU类描述。(默认 = [])
  13. # excludeUUIDexcludeID-除了它使用 UUID 之外,其他相同。(默认 = [])
  14. # 输出
  15. # deviceIDs - 所有可用 GPU 设备 ID 的列表。如果当前负载和内存使用量分别小于maxLoad和maxMemory,则认为 GPU 可用。该列表是根据 排序的order。返回的设备 ID 的最大数量由 限制limit。
  16. def get_gpu_ids(content):
  17. deviceIDs = GPUtil.getAvailable(order=content["gpu"]["order"],
  18. limit=int(content["gpu"]["limit"]),
  19. maxLoad=float(content["gpu"]["maxLoad"]),
  20. maxMemory=float(content["gpu"]["maxMemory"]),
  21. includeNan=content["gpu"]["includeNan"],
  22. excludeID=content["gpu"]["excludeID"],
  23. excludeUUID=content["gpu"]["excludeUUID"])
  24. return deviceIDs
  25. def get_all_gpu_ids():
  26. return GPUtil.getGPUs()