如何确定Deep Java Library (DJL)中的可用GPU内存？

Question

有没有办法在DJL中查询GPU上的最大总内存？例如，NVIDIA 1080ti应返回11 of的值。

我正在使用mxnet引擎，但与引擎无关的解决方案当然更可取。

Answer 1

是的，大疆已经通过https://github.com/awslabs/djl/blob/d2c47d0f2d663b8a5794c21d971420a99f2d47cd/api/src/main/java/ai/djl/util/cuda/CudaUtils.java中的JNA直接集成了CUDA API。所以这是一个与引擎无关的解决方案。

MemoryUsage mem = CudaUtils.getGpuMemory(device);
mem.getMax(); // it should return 11GB