WebGPU-Torch

基于 WebGPU 加速的张量计算

WebGPU-Torch 是一个受 PyTorch 启发的项目，它提供了基于 WebGPU 加速的张量计算和自动微分功能。这个项目的主页是 https://praeclarum.org/webgpu-torch。

安装和使用

WebGPU-Torch 没有依赖，所以你可以直接在你的网页中包含它，或者通过 npm 安装。在使用 WebGPU 张量之前，你需要先初始化库。初始化是一个异步函数，它会查询 navigator 对象以获取一个有效的 GPUDevice。

if (!await torch.initWebGPUAsync()) {
 console.warn(`WebGPU is not supported.`);
}

GPU 加速的张量

你可以创建一个张量，然后进行一些操作，比如加法。张量使用 WebGPU 内存（GPUBuffers）来存储它们的数据。当我们想要从张量中读取值时，我们必须将它映射到 CPU 地址空间。这可以通过 await a.storage.mapReadAsync() 或 await a.toArrayAsync() 来完成。

// Create a tensor
const a = torch.tensor([[1, 2, 3], [4, 5, 6]]);
// Create another tensor
const b = torch.tensor([[7, 8, 9], [10, 11, 12]]);
// Add them
const c = a.add(b);
const floatArray = await c.toArrayAsync();
console.log(floatArray);

自动微分支持

WebGPU-Torch 也支持自动微分。你可以创建一个张量，进行一些操作，然后进行微分。

// Create a tensor
const a = torch.tensor({data: [[1, 2, 3], [4, 5, 6]], requiresGrad: true});
// Create another tensor
const b = torch.tensor({data: [[7, 8, 9], [10, 11, 12]], requiresGrad: true});
// Add them
const c = a.add(b);
// Differentiate them
c.backward();

在这段代码执行后，a.grad、b.grad 和 c.grad 中将会有梯度张量值。

API

虽然这个库是受 PyTorch 启发的，但它并不是一个克隆，而是从零开始编写的。因此，它的 API 表面并不100%兼容 PyTorch，但作者尽可能地使它相似。

总的来说，WebGPU-Torch 是一个有趣的项目，它试图将深度学习带到 WebGPU 平台。虽然它目前还在开发中，但它的存在说明了社区对于支持更多硬件平台的需求。

本文作者：Maeiee

本文链接：WebGPU-Torch

版权声明：如无特别声明，本文即为原创文章，版权归 Maeiee 所有，未经允许不得转载！

喜欢我文章的朋友请随缘打赏，鼓励我创作更多更好的作品！