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

Posted 2023-03-16

【中文标题】加载了 tensorflow.js 的图模型可以在 GPU 上使用数据而不先将其传输到 CPU 吗？

【英文标题】：Can a graph model loaded with tensorflow.js use data on GPU without transferring it to the CPU first?

【发布时间】：2021-12-09 09:49:43

【问题描述】：

我目前正在使用 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 上？

【参考方案1】：

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

【参考方案2】：

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