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.gradb.gradc.grad 中将会有梯度张量值。

API

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

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


本文作者:Maeiee

本文链接:WebGPU-Torch

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


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