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

Question

I am currently using tfjs 3.8 to load a segmentation model (loaded as a tf.GraphModel ) on the client side. To create the input Tensor , I call browser.fromPixels(imageData) , which creates the Tensor on CPU from the ImageData object that is also on CPU. Since I'm using tfjs' webgl backend, the data is sent to the GPU when calling the model.predict(tensor) function. All of this works well, excepted that my ImageData object is created from an image on a canvas with a WebGLRenderingContext , meaning it comes from the GPU. This GPU->CPU->GPU data transfer is slowing down my process, which I am trying to optimize.

I briefly searched tfjs and could not find a way to create a Tensor on GPU to prevent the GPU->CPU data transfer. Is there a way I could keep my data on the GPU?

Answer 1

The solution is simply to provide the canvas with the webgl context to the browser.fromPixels(canvas) call. This will create the Tensor directly on the GPU.

Answer 2

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