webgpu-torch
v0.3.5
Published
Tensor computation and autograd with WebGPU acceleration
Downloads
26
Maintainers
Readme
webgpu-torch
Tensor computation and autograd with WebGPU acceleration inspired by pytorch.
Homepage: https://praeclarum.org/webgpu-torch
Installation
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
Usage
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.
GPU Accelerated Tensors
// 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);
Autograd Support
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
.
API
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.
Fundamental Types
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 adevice
, adtype
, ashape
, andstorage
properties. It can be created in a variety of ways.- Directly with
torch.tensor(array)
ornew torch.Tensor(array)
- From factory functions like
torch.zeros(shape)
ortorch.ones(shape)
- From operations like
a.add(b)
ora.mm(b)
- From a gradient calculation like
a.add(b).backward()
- Directly with
AutoFunction
is the base class for all autograd functions. It has aforward
method that computes the output tensor and abackward
method that computes the gradients of the inputs. They live in thetorch.functions
object. Functions should be called using theirapply
method.Kernel
is basic operation that can be executed on the GPU.
Tensor Operations
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.
TODO
Here are the big components of the library:
- [x] GPU Tensors
- [x] GPU Kernels
- [x] Autograd Functions
- [ ] Datatypes beyond float32
- [ ] Optimizers (SGD and Adam)
- [ ] Modules
- [ ] Save and restore (ONNX, safetensors)
In terms of supported operations, there's still a bit of work to be done:
- [x] Basic math
- [ ] Convolution
- [ ] Indexing
- [ ] Broadcasting
- [ ] Reductions
- [ ] Imaging
Acknowledgements
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.