Tensor computation and autograd with WebGPU acceleration inspired by pytorch.
Homepage: https://praeclarum.org/webgpu-torch
Webgpu-torch has no dependencies so you can just include it in your web page.
<script src="https://cdn.jsdelivr.net/npm/webgpu-torch@latest/torch.js"></script>
You can also include it using npm:
npm i webgpu-torch
If you want to use WebGPU tensors, you must first initialize the library with:
if (!await torch.initWebGPUAsync()) {
console.warn(`WebGPU is not supported.`);
}
It's an async function and will really do want to await it before doing anything else.
All it does is query the navigator
object for a valid GPUDevice
.
Sadly, that's an async operation.
// 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);
Tensors use WebGPU memory (GPUBuffers
) to store their data.
When we want to read values from the tensor we have to map it to the CPU address space.
This can be accomplished at a low level with await a.storage.mapReadAsync()
or at a high level with await a.toArrayAsync()
. Most functions in this library present a synchronous interface, but they are all asynchronous under the hood. Mapping the data to the CPU address space is the only visibly asynchronous operation in the library.
const floatArray = await c.toArrayAsync();
console.log(floatArray);
Math is fun, but it's even more fun when you do it backwards.
// 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();
After this code executes, there will be gradient tensor values in a.grad
, b.grad
, and c.grad
.
Although this library was inspired by pytorch, it is not a clone and was written from scratch. Its API surface is therefore not 100% compatible with pytorch, but I prioritize making it as similar as possible.
Device
is an abstraction over CPUs and GPUs allowing you to specify where tensors are allocated and executed.Dtype
is the data type of tensors and are specified as strings. Currently only "float32"
is supported.Shape
is an array of integers that specifies the size of each dimension of a tensor. For example, [32, 3, 128, 256]
would be 32 batched 256x128 RGB images.Tensor
is a multi-dimensional array of data. It has a device
, a dtype
, a shape
, and storage
properties. It can be created in a variety of ways.torch.tensor(array)
or new torch.Tensor(array)
torch.zeros(shape)
or torch.ones(shape)
a.add(b)
or a.mm(b)
a.add(b).backward()
AutoFunction
is the base class for all autograd functions. It has a forward
method that computes the output tensor and a backward
method that computes the gradients of the inputs. They live in the torch.functions
object. Functions should be called using their apply
method.Kernel
is basic operation that can be executed on the GPU.You have your basic unary operations like abs
that can be called from a global function or on the tensor directly:
const a = torch.tensor([[-1, 2, -3], [4, -5, 6]]);
const abs = torch.abs(a);
const abs2 = a.abs();
Your binary operations like add
can be called in the same way:
const b = torch.tensor([[7, -8, 9], [-10, 11, -12]]);
const sum = torch.add(a, b);
const sum2 = a.add(b);
I'm working on documenting the full list. For now, checkout the file op_table.ts for a list of most of the operations.
Here are the big components of the library:
In terms of supported operations, there's still a bit of work to be done:
I want to thank the Torch7 Lua environment for getting me into neural networks.
I want to thank the pytorch team for inspiring me.
I want to thank the webgpu teams at all the browser vendors for making this possible.
Generated using TypeDoc