Java Tensorflow SaveModelBundle 运行模型

首先需要有训练好的模型。

引入tensorflow的包到pom中

        <dependency>
            <groupId>org.tensorflow</groupId>
            <artifactId>tensorflow</artifactId>
            <version>1.15.0</version>
        </dependency>

ConfigProto config = ConfigProto.newBuilder()
                            .putDeviceCount("CPU", 8).setIntraOpParallelismThreads(50)
                            .setInterOpParallelismThreads(50)
                            .build();
// 增加模型运行配置

CofigProto在SaveModelBundle初始化时加载。源码：

    public Loader withConfigProto(byte[] configProto) {
      this.configProto = configProto;
      return this;
    }

创建SaveModelBundle实例化，构造模型对象。

SavedModelBundle.loader(modelUri).withTags("serve").withConfigProto(config.toByteArray()).load();

调用的native方法。

  private static native SavedModelBundle load(
      String exportDir, String[] tags, byte[] config, byte[] runOptions);

下面就是跑模型的方法了。

// feed为opName,参数 fetch为opName
public static float[][] run(SavedModelBundle savedModelBundle, Map<String, float[][]> feed, List<String> fetch) {
        float[][] r = new float[0][0];

        if (Objects.isNull(savedModelBundle)) {
            return r;
        }
        List<Tensor> tensorList = Lists.newArrayList();
        try {
            Session.Runner runner = savedModelBundle.session().runner();
            int size = 0;
            // 将feed opNmae和参数构造tensor，传入session中
            for (String feedOpName : feed.keySet()) {
                float[][] floats = feed.get(feedOpName);
                Tensor<?> inputTensor = Tensor.create(floats);
                tensorList.add(inputTensor);
                runner.feed(feedOpName, inputTensor);
                size = floats.length;
            }
            // fetch中opName传入session中
            for (String fetchOpName : fetch) {
                runner.fetch(fetchOpName);
            }
            // run模型
            List<Tensor<?>> out = runner.run();
            Tensor<?> tensor = out.get(0);
            tensorList.addAll(out);
            r = new float[size][1];
            tensor.copyTo(r);
        } catch (Exception ex) {
            logger.error(ex.getMessage(), ex);
        } finally {
            //结束时关闭tensor，tensor会占用内存，不关闭的话，容易出现oom。
           for (Tensor tensorClose : tensorList) {
               tensorClose.close();
           }
        }
        return r;
    }

Tensor源码：

/**
   * Create a Tensor of data type {@code dtype} from a Java object. Requires the parameter {@code T}
   * to match {@code type}, but this condition is not checked.
   *
   * @param obj the object supplying the tensor data.
   * @param dtype the data type of the tensor to create. It must be compatible with the run-time
   *     type of the object.
   * @return the new tensor
   */
  private static Tensor<?> create(Object obj, DataType dtype) {
    @SuppressWarnings("rawtypes")
    Tensor<?> t = new Tensor(dtype);
    t.shapeCopy = new long[numDimensions(obj, dtype)];
    fillShape(obj, 0, t.shapeCopy);
    long nativeHandle;
    if (t.dtype != DataType.STRING) {
      int byteSize = elemByteSize(t.dtype) * numElements(t.shapeCopy);
      nativeHandle = allocate(t.dtype.c(), t.shapeCopy, byteSize);
      setValue(nativeHandle, obj);
    } else if (t.shapeCopy.length != 0) {
      nativeHandle = allocateNonScalarBytes(t.shapeCopy, (Object[]) obj);
    } else {
      nativeHandle = allocateScalarBytes((byte[]) obj);
    }
    t.nativeRef = new NativeReference(nativeHandle);
    return t;
  }

session run 的源码：

private Run runHelper(boolean wantMetadata) {
      long[] inputTensorHandles = new long[inputTensors.size()];
      long[] inputOpHandles = new long[inputs.size()];
      int[] inputOpIndices = new int[inputs.size()];
      long[] outputOpHandles = new long[outputs.size()];
      int[] outputOpIndices = new int[outputs.size()];
      long[] targetOpHandles = new long[targets.size()];
      long[] outputTensorHandles = new long[outputs.size()];

      // It's okay to use Operation.getUnsafeNativeHandle() here since the safety depends on the
      // validity of the Graph and graphRef ensures that.
      int idx = 0;
      for (Tensor<?> t : inputTensors) {
        inputTensorHandles[idx++] = t.getNativeHandle();
      }
      idx = 0;
      for (Output<?> o : inputs) {
        inputOpHandles[idx] = o.getUnsafeNativeHandle();
        inputOpIndices[idx] = o.index();
        idx++;
      }
      idx = 0;
      for (Output<?> o : outputs) {
        outputOpHandles[idx] = o.getUnsafeNativeHandle();
        outputOpIndices[idx] = o.index();
        idx++;
      }
      idx = 0;
      for (GraphOperation op : targets) {
        targetOpHandles[idx++] = op.getUnsafeNativeHandle();
      }
      Reference runRef = new Reference();
      byte[] metadata = null;
      try {
        metadata =
            Session.run(
                nativeHandle,
                runOptions,
                inputTensorHandles,
                inputOpHandles,
                inputOpIndices,
                outputOpHandles,
                outputOpIndices,
                targetOpHandles,
                wantMetadata,
                outputTensorHandles);
      } finally {
        runRef.close();
      }
      List<Tensor<?>> outputs = new ArrayList<Tensor<?>>();
      for (long h : outputTensorHandles) {
        try {
          outputs.add(Tensor.fromHandle(h));
        } catch (Exception e) {
          for (Tensor<?> t : outputs) {
            t.close();
          }
          outputs.clear();
          throw e;
        }
      }
      Run ret = new Run();
      ret.outputs = outputs;
      ret.metadata = metadata;
      return ret;
    }

 private static native byte[] run(
      long handle,
      byte[] runOptions,
      long[] inputTensorHandles,
      long[] inputOpHandles,
      int[] inputOpIndices,
      long[] outputOpHandles,
      int[] outputOpIndices,
      long[] targetOpHandles,
      boolean wantRunMetadata,
      long[] outputTensorHandles);
}