如何在DJL (Deep Java Library)中调用自定义mxnet运算符？

Question

如何从DJL调用自定义mxnet运算符？例如来自examples的my_gemm运算符。

Answer 1

可以像内置的mxnet引擎一样手动调用JnaUtils，只需使用定制的库即可。对于my_gemm示例，如下所示：

import ai.djl.Device;
import ai.djl.mxnet.jna.FunctionInfo;
import ai.djl.mxnet.jna.JnaUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;
import ai.djl.util.PairList;

import java.util.Map;

// Load the external mxnet operator library
JnaUtils.loadLib("path/to/incubator-mxnet/example/extensions/lib_custom_op/libgemm_lib.so", 1);
// get a handle to the loaded operator
Map<String, FunctionInfo> allFunctionsAfterLoading = JnaUtils.getNdArrayFunctions();
FunctionInfo myGemmFunction = allFunctionsAfterLoading.get("my_gemm");
// create a manager to execute the example with
try (NDManager ndManager = NDManager.newBaseManager().newSubManager(Device.cpu())) {
    // create input for the gemm call
    NDArray a = ndManager.create(new float[][]{new float[]{1, 2, 3}, new float[]{4, 5, 6}});
    NDArray b = ndManager.create(new float[][]{new float[]{7}, new float[]{8}, new float[]{9}});
    // call the function manually (NDManager.invoke will not work, as it caches the mxnet
    // engine operators and ignores external ones)
    PairList<String, Object> params = new PairList<>();
    NDArray result = myGemmFunction.invoke(ndManager, new NDArray[]{a, b}, params)[0];
    // prints
    // ND: (2, 1) cpu() float32
    //[[ 50.],
    // [122.],
    //]
    // (same as the python example)
    System.out.println(result);
}