如何在运行时确定dl4j / nd4j后端是CPU还是GPU?

时间:2018-03-06 18:51:17

标签: kotlin dl4j nd4j

dl4j的优化仅适用于GPU: DataTypeUtil.setDTypeForContext(DataBuffer.Type.HALF)

如果后端是GPU,我只想拨打电话。

在我的Maven pom.xml中,我已经

<!-- CPU or GPU -->
<nd4j.backend>nd4j-native-platform</nd4j.backend>
<!--<nd4j.backend>nd4j-cuda-8.0-platform</nd4j.backend>-->

我正在寻找从Java中读取这些价值的方法,所有这些看起来都很笨拙。如果我可以查询dl4j或nd4j以获得&#34;我们正在运行什么样的后端会更容易?&#34;然后根据它进行优化调用。

从回答编辑:

Nd4jBackend.load().let { be->
    println("nd4j Backend: ${be.javaClass.simpleName}")
    if(be.javaClass.simpleName.toLowerCase().contains("gpu")) {
        println("Optimizing for GPU")
        DataTypeUtil.setDTypeForContext(DataBuffer.Type.HALF)
    }
}

2 个答案:

答案 0 :(得分:1)

看看您是否可以使用Nd4j.backend。使用cuda打印它我得到:

org.nd4j.linalg.jcublas.JCublasBackend

并且没有cuda:

org.nd4j.linalg.cpu.nativecpu.CpuBackend

答案 1 :(得分:0)

当你启动nd4j时,它也会在开头打印出来。应该有一个供应商为后端打印出来。