java - 使用 TensorFlow for Java 的内存泄漏

Question

以下测试代码会泄漏内存:

private static final float[] X = new float[]{1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0};

public void testTensorFlowMemory() {
    // create a graph and session
    try (Graph g = new Graph(); Session s = new Session(g)) {
        // create a placeholder x and a const for the dimension to do a cumulative sum along
        Output x = g.opBuilder("Placeholder", "x").setAttr("dtype", DataType.FLOAT).build().output(0);
        Output dims = g.opBuilder("Const", "dims").setAttr("dtype", DataType.INT32).setAttr("value", Tensor.create(0)).build().output(0);
        Output y = g.opBuilder("Cumsum", "y").addInput(x).addInput(dims).build().output(0);
        // loop a bunch to test memory usage
        for (int i=0; i<10000000; i++){
            // create a tensor from X
            Tensor tx = Tensor.create(X);
            // run the graph and fetch the resulting y tensor
            Tensor ty = s.runner().feed("x", tx).fetch("y").run().get(0);
            // close the tensors to release their resources
            tx.close();
            ty.close();
        }

        System.out.println("non-threaded test finished");
    }
}

有什么明显的我做错了吗？基本流程是创建一个图和该图上的一个 session ，创建一个占位符和一个常量，以便对作为 x 输入的张量进行累加。运行生成的 y 操作后，我关闭 x 和 y 张量以释放它们的内存资源。

到目前为止我认为有帮助的事情:

• 这不是 Java 对象内存问题。根据 jvisualvm，堆不会增长，JVM 中的其他内存也不会增长。根据 Java 的 native 内存跟踪，这似乎不是 JVM 内存泄漏。
• 关闭操作很有帮助，如果它们不存在，内存会突飞猛进地增长。有了它们，它的生长速度仍然非常快，但几乎和没有它们一样快。
• cumsum 运算符并不重要，sum 和其他运算符也会出现这种情况
• 它发生在带有 TF 1.1 的 Mac OS 和带有 TF 1.1 和 1.2_rc0 的 CentOS 7 上
• 注释掉 Tensor ty 行可以消除泄漏，因此它看起来就在那里。

有什么想法吗？谢谢!另外，here's a Github project that demonstrates this issue同时进行线程测试(以更快地增加内存)和非线程测试(以表明它不是由于线程)。它使用 maven 并且可以简单地运行:

mvn test

Answer 1

我相信确实存在泄漏(特别是缺少与 allocation in JNI code 对应的 TF_DeleteStatus)(感谢详细的重现说明)

我鼓励您在 http://github.com/tensorflow/tensorflow/issues 提交问题 并希望它应该在最终的 1.2 版本之前得到修复。

(相关地，由于 Tensor.create(0) 创建的 Tensor 对象未被关闭，因此您在循环外也有泄漏)

更新:已修复，1.2.0-rc1 应该不会再有这个问题。