Tinygrad 微型DL框架

显然,Tinygrad似乎只是另一个深度学习框架。 嗯,不完全是。 它要比Pytorch更小和更简单几个数量级。 每个曾经尝试阅读 PyTorch 或 TensorFlow 底层源代码的人都会对此表示赞赏。 它无疑保证了添加更多加速器的门槛足够低。

由于tinygrad极其简单,它的目标是成为添加新加速器的最简单的框架,并支持推理和训练。

这个名字从何而来?

tinygrad 将始终低于 1000 行。 如果不是,我们将恢复提交,直到tinygrad变得更小。

我认为这有点过于雄心勃勃,除非你将其分成各种包来添加各种支持,例如各种加速器、模型结构、配置文件和加载器等等。 但是,是的,这听起来确实很有趣。

加速的内容已经移入 Accel

1、tinygrad代码结构

pyreverse --filter-mode ALL --colorized --max-color-depth 7 --output webp ./tinygrad/

输出以下代码结构:

2、加速器的使用

  • GPU

NVIDIA 的 CUDA 无疑在神经网络训练领域占据主导地位。 因此,看到加速器支持使用 OpenCL 是一个很好的改变。 那么,这真的是 GPU 特定的吗? 不,某些 CPU 也支持 OpenCL 的某些子集(例如通过 pocl)。 使用 OpenCL 的有趣之处在于支持完全不同的硬件来加速训练或推理,例如 FPGA。 如果我没记错的话,英特尔的一些 FPGA 可以使用 OpenCL 进行编程。 然而,我不确定 JIT 编译的 PyOpenCL 内核如何容易地转换为编程的输入,例如 FPGA(如果它们不提供任何 OpenCL 功能)。

然而,似乎添加了通过 PyCUDA 的 CUDA 支持,但尚未集成。

  • ANE(苹果神经引擎)

张量类似乎支持一些 ANE 操作。 他们似乎正在构建自己的 ANE 抽象层,以利用 ANE 进行培训。

3、tinygrad上手试用

注意! 如果我们在 conda 环境中使用tinygrad,那么我们必须安装 OpenCL 实现(例如 pocl)

conda install -c conda-forge pocl pyopencl clinfo

clinfo 对于列出可用的 OpenCL 设备很有用,但它需要安装在 conda 环境中。

如果我们使用标准示例,我们可以看到有 CPUBuffer 和 GPUBuffer:

from tinygrad.tensor import Tensor
(Tensor.ones(5000,5000).gpu() + Tensor.ones(5000,5000).gpu()).cpu()

输出如下:

<Tensor CPUBuffer([[2., 2., 2., ..., 2., 2., 2.],
           [2., 2., 2., ..., 2., 2., 2.],
           [2., 2., 2., ..., 2., 2., 2.],
           ...,
           [2., 2., 2., ..., 2., 2., 2.],
           [2., 2., 2., ..., 2., 2., 2.],
           [2., 2., 2., ..., 2., 2., 2.]], dtype=float32) with grad None>

然而:

from tinygrad.tensor import Tensor
Tensor.ones(5000,5000).gpu() + Tensor.ones(5000,5000).gpu())

输出:

<Tensor <GPUBuffer with shape (5000, 5000)> with grad None>

如果导出 PYOPENCL_COMPILER_OUTPUT=1,PyOpenCL 将显示编译器警告:

CompilerWarning: Built kernel retrieved from cache. Original from-source build had warnings:
Build on <pyopencl.Device 'Intel(R) Iris(R) Xe Graphics [0x9a49]' on 'Intel(R) OpenCL HD Graphics' at 0x562deb03ccb0> succeeded, but said:

1:1:112: warning: double precision constant requires cl_khr_fp64, casting to single precision
inline float get_A(__global const float *x, int gid) { int valid = 1; int idx = gid; ; return valid ? x[idx] : 0.0;}
                                                                                                               ^
1:2:112: warning: double precision constant requires cl_khr_fp64, casting to single precision
inline float get_B(__global const float *x, int gid) { int valid = 1; int idx = gid; ; return valid ? x[idx] : 0.0;}inline float _ewop(int gid,float acc,__global const float *A_g,__global const float *B_g) {float A = get_A(A_g, gid);
                                                                                                               ^
1:4:19: warning: double precision constant requires cl_khr_fp64, casting to single precision
      float acc = 0.0;
                  ^

  warn(text, CompilerWarning)

这是一个快速 DNN 推理示例:

ipython3 examples/efficientnet.py https://media.istockphoto.com/photos/hen-picture-id831791190

看来 export GPU=1 终于可以体验到一些加速了。

使用 Intel i7-1165G7 的结果:

8 8.045593 hen
did inference in 1.10 s

使用 GPU 加速:

8 8.0455885 hen
did inference in 0.31 s

4、结束语

Tinygrad 确实还不够成熟。 然而,我想最好让每个人自己来决定tinygrad是一个笑话还是在一个过于复杂的深度学习框架的世界中迫切需要的东西。


原文链接:Introducing tinygrad

BimAnt翻译整理,转载请标明出处