typio / nervous

Deep Learning Framework in WebGPU (For Education)

nervous

Web Machine Learning Framework

Docs • Demos

Overview

An ML framework for the browser, with these goals:

• simple yet efficient
• less convoluted than TensorFlow and pyTorch
• supporting levels of interaction from tensor operations to full neural nets.

Features GPU acceleration using WebGPU and compute shaders.

I decided to remove the dual-backend system supporting JavaScript and WebGPU and go forward with just WebGPU. It's a hard choice because I've invested so much in the JS backend, but it will be useless outside of educational purposes once WebGPU is released and removing the multiple backend support will reduce the architectural complexity of the library and the lag time of new features.

Usage

package is not yet published

import nv from "@typio/nervous";

const main = async () => {
    await nv.init();
    const [n, m] = await Promise.all([
        nv.random([2047, 512]),
        nv.random([512, 2047]),
    ]);
    let res = n.dot(m);

    await res.print(4);
};

main();

nervous's Issues

Should functions return a Promise or be resolved internally?

Solve WebGPU OP exceeding max compute workgroups

Potential Solutions

1. Settle with current limits
(only guarantee <=2047x2047, result of workgroup_size(64))

1. Utilize x,y maybe z at 64
(cons: verbose wgsl with conversions from 2d/3d to 1d indicies)

2. Loops in wgsl to extend
(cons: technically lessens parallelization, how much would be enough??)

3. Chunk tensor op?
(cons: buffer memory will overflow anyway and this would add a much more complexity)

Consider removing JS backend along with multiple backend concept