tinygrad/tinygrad

tinygrad 是一个端到端的深度学习技术栈： - 带有 autograd 的 **Tensor 库** - 融合并底层化 kernel 的 **IR 和编译器** - **JIT + 图执行** - 用于真实训练的 **nn / optim / datasets** ## tinygrad 的对比 **PyTorch** - ✅ 相似：eager `Tensor` API、autograd、`optim`、基础数据集和层。 - ✅ 你可以编写熟悉的训练循环。 - 🔁 与 PyTorch 不同，整个编译器和 IR 都是可见且可修改的。 **JAX** - ✅ 基于 primitives 的 IR 自动微分（类似于 JAXPR + XLA）。 - ✅ 函数级 JIT (`TinyJit`)，可以捕获并重放 kernel。 - 🔁 函数式转换较少（目前还没有完整的 `vmap`/`pmap`），但可读性强得多。 **TVM** - ✅ 对 kernel 进行多次 lowering pass、调度和 BEAM 搜索。 - ✅ 用于批量执行的设备“图”。 - 🔁 tinygrad 还提供了**前端框架**（tensors、nn、optim），而不仅仅是编译器。 ### 惰性 尝试一个 matmul。看看它是如何凭借惰性的力量，被融合到一个 kernel 中的。 ``` DEBUG=3 python3 -c "from tinygrad import Tensor; N = 1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N); (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()" ``` 我们可以将 `DEBUG` 更改为 `4` 来查看生成的代码。 ### 神经网络 事实证明，对于神经网络来说，你需要的 90% 都是一个不错的 autograd/tensor 库。 加上一个 optimizer、一个 data loader 和一些计算，你就拥有了所需的一切。 ``` from tinygrad import Tensor, nn, Context class LinearNet: def __init__(self): self.l1 = Tensor.kaiming_uniform(784, 128) self.l2 = Tensor.kaiming_uniform(128, 10) def __call__(self, x:Tensor) -> Tensor: return x.flatten(1).dot(self.l1).relu().dot(self.l2) model = LinearNet() optim = nn.optim.Adam([model.l1, model.l2], lr=0.001) x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7]) # replace with real mnist dataloader with Context(TRAINING=1): for i in range(10): optim.zero_grad() loss = model(x).sparse_categorical_crossentropy(y).backward() optim.step() print(i, loss.item()) ``` 查看 [examples/beautiful_mnist.py](examples/beautiful_mnist.py) 获取在约 5 秒内达到 98% 准确率的完整版本 ## 加速器 tinygrad 已经支持众多加速器，包括： - [x] [OpenCL](tinygrad/runtime/ops_cl.py) - [x] [CPU](tinygrad/runtime/ops_cpu.py) - [x] [METAL](tinygrad/runtime/ops_metal.py) - [x] [CUDA](tinygrad/runtime/ops_cuda.py) - [x] [AMD](tinygrad/runtime/ops_amd.py) - [x] [NV](tinygrad/runtime/ops_nv.py) - [x] [QCOM](tinygrad/runtime/ops_qcom.py) - [x] [WEBGPU](tinygrad/runtime/ops_webgpu.py) 运行以下命令检查默认加速器：`python3 -c "from tinygrad import Device; print(Device.DEFAULT)"` ## 安装 目前推荐的安装 tinygrad 的方式是从源码安装。 ### 从源码安装 ``` git clone https://github.com/tinygrad/tinygrad.git cd tinygrad python3 -m pip install -e . ``` ### 直接安装 (master) ``` python3 -m pip install git+https://github.com/tinygrad/tinygrad.git ``` ## 文档 文档以及快速入门指南可以在由 [docs/](/docs) 目录构建的 [文档网站](https://docs.tinygrad.org/) 上找到。 ### 与 PyTorch 对比的快速示例 ``` from tinygrad import Tensor x = Tensor.eye(3) y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x).sum() z.backward() print(x.grad.tolist()) # dz/dx print(y.grad.tolist()) # dz/dy ``` 在 PyTorch 中的相同实现： ``` import torch x = torch.eye(3, requires_grad=True) y = torch.tensor([[2.0,0,-2.0]], requires_grad=True) z = y.matmul(x).sum() z.backward() print(x.grad.tolist()) # dz/dx print(y.grad.tolist()) # dz/dy ``` ### 运行测试 你应该使用 `pre-commit install` 安装 pre-commit hooks。这会在每次提交时运行 linter、mypy 和一部分测试。 有关如何运行完整测试套件的更多示例，请参阅 [CI workflow](.github/workflows/test.yml)。 在本地运行测试的一些示例： ``` python3 -m pip install -e '.[testing]' # install extra deps for testing python3 test/backend/test_ops.py # just the ops tests python3 -m pytest test/ # whole test suite ``` #### Process replay 测试 [Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/README.md) 会将你的 PR 生成的 kernel 与 master 分支进行对比。如果你的 PR 是一个重构或提速，且没有任何预期的行为改变，则应在 pull request 标题中包含 [pr]。

标签：Vectored Exception Handling, Veh, 逆向工具