加载 tensorflow.js 的图形模型是否可以在不先将数据传输到 CPU 的情况下在 GPU 上使用数据?

Can a graph model loaded with tensorflow.js use data on GPU without transferring it to the CPU first?

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

我简单地搜索了 tfjs,找不到在 GPU 上创建一个 Tensor 来阻止 GPU->CPU 数据传输的方法。有什么方法可以将我的数据保存在 GPU 上?

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

有关此主题的详细对话在线程中 https://github.com/tensorflow/tfjs/issues/5765