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

Question

当前在 Github 上发布的 Tensorflow.js 包括一个 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() 才能对其进行任何操作。

希望对您有所帮助。