如何在 Tensorflow.JS 上激活 webgpu 后端

Question

Tensorflow.js Github 上的当前版本包括一个 WEBGPU 后端，但是当运行 tf.getBackend() 我得到 webgl 即使 Chrome启用了 WEBGPU 的 Canary。另外运行 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
})()

如何在 Tensorflow.JS 上激活 webgpu 后端

How to activate webgpu backend on Tensorflow.JS

tensorflow.js