前言
AI开发通常使用到GPU,但通常使用的是python、c++等语言,java用的则非常少。这也导致了java在gpu相关的库比较少。现在的需求是要获取nvidia显卡的使用情况,如剩余显存。这里给出两种较简单的解决方案。
基于nivdia-smi工具
显卡是硬件,当然需要操作系统支持相关的驱动程序。首先要确保的正确安装了nvidia的相关驱动程序。输入命令nvidis-smi
后有类似下面的输出,即证明当前机器已安装好了驱动程序。我们这个方案也是基于nvidia-smi
这个命令来做的。
我们可以看到上图的输出,表面了当前的驱动版本、CUDA版本、显卡数量、显存、显卡利用率,当前使用进程等信息。但是明显是给人来看的,我们想放到程序中则需要进行解析。好在该工具提供了xml格式输出:
于是有了思路:
- 通过进程调用
nvidia-smi
,并获取xml输出结果 - 解析xml格式,获取需要的内容
- 将需要的内容转为java bean以供使用
具体的xml格式使用"nvidia-smi -q -x"
一看便知,我在这里也就不赘述如何解析了。这里给出一个简单的库,目前只包含了显存信息的解析,其余内容没有实现:https://gitee.com/angelhand/vel-koz
基于nvml绑定库
上面提到的nvidia-smi
工具实际上是英伟达基于工具nvml实现的,这里有对该工具的一些介绍:https://developer.nvidia.com/nvidia-management-library-nvml。
该工具是一个c编写的管理工具,上面文档有介绍具体的API。工具可以通过GPU Deployment Kit一并安装,GPU Deployment Kit又是CUDA Toolkit的一部分。不过CUDA8.0之后,GPU Deployment Kit不再单独提供,而是合并到了CUDA Toolkit里了。什么意思呢,就是CUDA8.0之后,要使用nvml,只需要安装CUDA Toolkit:https://developer.nvidia.com/cuda-downloads。使用这个方案是一定要安装的哦。不过Windows11环境下,使用系统更新自动安装好的驱动就可以了,不需要额外下载(40系显卡31版本驱动,旧版本不保证)。
环境安装好之后使用https://github.com/henkelmax/nvmlj这个绑定库。这里有同学要问绑定库是什么意思,简单解释一下。前面提到的nvml工具是c写的,提供了一个nvml.h
的头文件,其他语言要使用就需要对这些功能做一个映射(绑定)。这个头文件就类似于java中的接口。java对c的映射可以使用JNI技术,另外还有比较简单的JNA。我们这个库就是基于JNA实现的。下面来展示如何使用。
首先修改一下源:
<repository><id>henkelmax.public</id><url>https://maven.maxhenkel.de/repository/public</url>
</repository>
然后使用进行依赖引入:
<dependency><groupId>de.bommel24.nvmlj</groupId><artifactId>nvmlj</artifactId><version>1.0.2</version>
</dependency>
使用起来也很简单:
public class Main {public static void main(String[] args) throws NVMLJException {// The path to the nvml library (Defaults to C:/Program Files/NVIDIA Corporation/NVSMI/nvml.dll)// windows系统需要这样设置System.setProperty("nvml.path", "C:\\Windows\\System32\\nvml.dll");// 这是我的wsl找到的,实体机可以whereis一下System.setProperty("nvml.path", "/usr/lib/wsl/lib/libnvidia-ml.so.1");NVMLJ.nvmlInit();System.out.println(NVMLJ.nvmlDeviceGetCount());// 不要忘记关闭资源NVMLJ.nvmlShutdown();}
}
这里有个点需要注意一下,就是代码里那行System.setProperty()
。set的内容实际上就是nvml的动态链接库。Windows较新版本的目录就是上面提到的目录,可以找一下是否存在。Linux下,我使用的是wsl ubuntu22.04系统,在上面那个目录里。如果你使用的是其他版本的系统,可以whereis libnvidia-mi.so.1
来查找一下库的具体位置。你可能会查出来多个位置,或者没有。那么再查一下libnvidia-ml.so
(后面没有.1
)。具体哪个管用试试就知道了,不过通常是会出现再usr
的lib
/lib32
/lib64
的子目录或者子子(子…)目录下了。