如何在Tensorflow.JS上开通webgpu后台

Question

在Tensorflow.js的Github上的当前版本包括一个WEBGPU后端，但当运行tf.getBackend()时，即使启用了WEBGPU的Chrome Canary，我也得到了webgl。此外，运行tf.backend()似乎并不意味着WEBGPU就在那里。

Answer 1

您要导入哪个tf包？默认情况下，WebGPU后端不与tfjs捆绑在一起-您需要专门导入它，如下所示：

import * as tf from '@tensorflow/tfjs-backend-webgpu';

const init = async () => {
    await tf.ready();

    // Now we can create tensors and run ops.
    tf.matMul(a, b).print(); 
};

init();

请注意，后端是异步的-在对其执行任何操作之前，您必须调用await tf.ready()。

希望这能有所帮助。

Answer 2

这仍然是一个实验性的特征。您可以使用npm导入包

 npm install @tensorflow/tfjs-backend-webgpu

并导入它

import * as tf from '@tensorflow/tfjs-backend-webgpu';

唯一要记住的就是等待后端

import * as tf
(async() => {
  await tf.ready 
 // then do all operations on the backend
})()