加载了 tensorflow.js 的图形模型可以使用 GPU 上的数据而不先将其传输到 CPU 吗？

Question

我目前正在使用 tfjs 3.8 在客户端加载分段模型（加载为tf.GraphModel ）。 要创建输入Tensor ，我调用browser.fromPixels(imageData) ，它从 CPU 上的ImageData对象创建 CPU 上的Tensor 。 由于我使用的是 tfjs 的webgl后端，因此在调用model.predict(tensor)函数时将数据发送到 GPU。 所有这些都运行良好，除了我的ImageData对象是从带有WebGLRenderingContext的画布上的图像创建的，这意味着它来自 GPU。 这个 GPU->CPU->GPU 数据传输减慢了我正在尝试优化的进程。

我简单地搜索了 tfjs 并找不到在 GPU 上创建Tensor以防止 GPU->CPU 数据传输的方法。 有没有办法将我的数据保存在 GPU 上？

Answer 1

解决方案只是为画布提供 webgl 上下文给browser.fromPixels(canvas)调用。 这将直接在 GPU 上创建张量。

Answer 2

关于这个主题的详细对话在一个线程https://github.com/tensorflow/tfjs/issues/5765