Taichi「量化」功能,到底能节省多少显存?

您好,欢迎访问我们的网站,我们将竭诚为您服务!

Taichi「量化」功能,到底能节省多少显存?

时间:2022-09-30 21:55:32 阅读:94
Taichi v1.1首次正式发布了量化数据类型(quantized data types)功能。在这篇文章中,我们将带大家一探「量化的威力」。

  什么是量化?

  量化(quantization)是数字信号处理中的一个概念,指的是把一个更大的集合以某种方式映射到一个更小的集合。例如,你可以把实数信号全部四舍五入成整数信号,这样需要表示的数的数量就大大减少了,可以以很小的空间存储在计算机中。当然,和原始信号相比,这样的做法也丢失了一定的精确度。

  但丢失精确度不代表结果是没有意义的。例如对于图像处理领域,只要做到产生的差异较小,人眼难以分辨,就完全可以采用精确度更低的结果。近年来,量化在神经网络等领域更是已经成为了提升计算效率的常用技巧。

  对于Taichi最常被使用的场景——物理仿真来说,少量精确度的牺牲往往也是可接受的。更好的视觉效果通常来自更高的分辨率,而有限的GPU显存往往会成为持续提升分辨率的障碍。在这种情况下,拥有一个能节省显存开销的方法就变得至关重要——Taichi的量化数据类型应运而生。借助量化数据类型,你可以在不改变计算部分的代码的前提下,将数据以任意低的精确度进行存储。在一些例子中,使用Taichi量化数据类型甚至可以节约高达8倍的显存。

  动图封面

  4亿体素的烟雾仿真,占用29G显存(量化后节约了1.6倍)

  动图封面

  2亿粒子的弹性体仿真,占用17G显存(量化后节约了1.7倍)

  你的第一个「量化」Taichi程序

  看到这里,想必你已经跃跃欲试,想知道怎么用量化数据类型去降低Taichi程序的显存占用了。在这一节里,我们将一起尝试魔改Taichi自带的一份示例代码,仅修改数据定义部分就让它迅速「瘦身」。我们选取的例子是examples/simulation/euler.py

  动图封面

  数据定义里大部分都是如下所示的field:

  Q=ti.Vector.field(4,dtype=ti.f32,shape=(N,N))

  这里的每个元素都使用ti.f32,也就是单精度浮点数类型来存储。假如我们能适当降低每个元素的精度,让一个向量的4个元素合计占用的二进制位数从128变为64,就可以让显存占用减半。

  这时候,Taichi的量化浮点数类型就闪亮登场了:

  float_8_8=ti.types.quant.float(exp=8,frac=8)

  我们知道,浮点数由指数、尾数和符号位构成。例如,单精度浮点数包含8位指数、23位尾数和1位符号,这是由IEEE 754标准规定的。而使用Taichi的量化浮点数类型,你就可以自定义指数和尾数的位数(为方便处理,Taichi将符号位包含在了尾数中)。在上面一行代码中,我们就定义出了一个包含8位指数(exp)和8位尾数(frac)的浮点数类型float_8_8。如果我们能把每个元素的存储类型替换成float_8_8,就可以达成上面所说的显存占用减半的目的。

  由于量化数据类型不是原生数据类型,不能单独存储,你需要按自己的需求把若干个量化数据类型一字排开放进某个原生数据类型。为此,Taichi引入了ti.BitpackedFields这个API,专门用于把一组量化数据类型的field捆在一起形成AOS(Array-of-Structure)排布,即这几个量化数据类型会合成一个原生数据类型进行存储。指定完自身的组成后,ti.BitpackedFields可以被放置在任意SNode下,从而放置其中的每个field:

  Q=ti.Vector.field(4,dtype=float_8_8)

  bitpack=ti.BitpackedFields(max_num_bits=64)

  bitpack.place(Q)

  ti.root.dense(ti.ij,(N,N)).place(bitpack)

  到这里,我们对数据定义的改动就全部结束了。让我们跑一下试试:

  动图封面

  左图为原版,右图为第一次量化修改后

  很不幸,仿真结果出现了比较大的偏差。我们虽然保留了和原来单精度浮点数相同的范围(指数不变),但是牺牲了过多的精确度(尾数从24减少到了8)。是不是我们至此就宣告失败了呢?别急,Taichi还有下一招!考虑到这些数都拥有实际物理意义,很多联系紧密的物理量(比如一个向量的各轴分量)可以共享指数(使用shared_exponent指定),从而给尾数留出更多空间(在这个例子里可以让尾数增加到14):

  float_8_14=ti.types.quant.float(exp=8,frac=14)

  Q=ti.Vector.field(4,dtype=float_8_14)

  bitpack=ti.BitpackedFields(max_num_bits=64)

  bitpack.place(Q,shared_exponent=True)

  ti.root.dense(ti.ij,(N,N)).place(bitpack)

  我们再跑一下试试,这次就很难分辨出和原版的差别了:

  动图封面

  左图为原版,右图为第二次量化修改后

  至此,你的第一个「量化」taichi程序宣告完成,并且在对视觉效果影响甚微的情况下节省了一半显存!你可能也体会到了,能在不改变计算部分的代码的前提下去修改、尝试不同的用于存储的数据类型有多么方便。这里我们只展示了量化数据类型完整功能的冰山一角,关于各种量化数据类型的详细介绍、更多示例和高级特性可以在我们的文档中找到。taichi https://taichi-lang.cn/
郑重声明:文章内容来自互联网,纯属作者个人观点,仅供参考,并不代表本站立场 ,版权归原作者所有!

上一篇:与 CUB、CuPy 和 Numba 相比,用 Taichi 做数值计算性能怎么样?

下一篇:什么是VPS远程桌面服务器?向日葵如何进行远程桌面?

相关推荐

返回顶部