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

Question

Github 上当前版本的 Tensorflow.js 包含 WEBGPU 后端，但运行时tf.getBackend()我webgl什至可以使用启用了 WEBGPU 的 Chrome 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()在对其执行任何操作之前调用。

希望有帮助。