Taichi & Torch 进阶篇详解二者的数据容器

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

Taichi & Torch 进阶篇详解二者的数据容器

时间:2022-10-23 23:01:51 阅读:89
在图形学领域,工程师经常遇到需要使用多个分量来描述一个属性的情形,对于熟悉如 GLSL 等着色语言(shading language)的同学可能对这个概念并不陌生。比如可以通过定义 RGBA 四个分量来描述一个颜色,或者通过定义 xyz 坐标描述三维空间的一个位置。因此,Taichi 的数据容器 field 中的元素既可以是标量,也可以是向量或矩阵,以便于灵活地适配图形学的各种应用场景。以向量元素为例,对于它的访问既可以使用数组下标,如 field[0] ,也可以使用如 field.r, field.g, field.x, field.y 等 Swizzle 操作。

举个例子,要表示一张 512 x 512 大小的图片中每个像素点的颜色 (RGB),你可以在 Taichi 中创建一个由三维向量组成的 field:

ti_image = ti.Vector.field(n=3, shape=(512, 512), dtype=ti.f32)

这样你就可以通过 x[10, 20] 访问到横坐标为 10, 纵坐标为 20 的像素点的颜色,并且可以通过 x[10, 20].r 来访问到该像素 R 通道的值。这样用户可以用一种符合直觉的方式便捷地进行位置、向量或者颜色向量的数学运算。

p.s. Taichi 的向量和矩阵元素会在编译时被展开。为避免展开影响整体性能,Taichi field 只支持小的向量和矩阵 如 vec3, mat4x4 等。

关于 Torch

在机器学习领域使用频繁的是基于 Tensor 的如矩阵乘法等操作,因此 torch.Tensor 的元素只允许是标量(包括复数),这样更符合数学上的张量定义。当然如果 Tensor 的某些维度有对应的物理含义,就需要用户来记住这些信息。

比如在 PyTorch 中,你可以创建一个 tensor 表示和上面 field ti_image
类似的一张 512 x 512 大小的图片中每个像素点的颜色,常用的做法是创建下面这样一个三维张量,由最后一维的三个值来隐式代表 RGB 三个通道的值。

torch_image = torch.zeros((height, width, 3), dtype=torch.float)


Taichi 与 Torch 的内存读取方式

除了定义数据的接口不同,Taichi 和 PyTorch 都支持用户灵活调整 field 或 Tensor 的内存排布,从而优化性能。

Pytorch Tensor 默认的内存排布是 NCHW,可以通过指定 memory_format 改变内存排布。机器学习应用中常用的两种格式包括:

「NHWC」的排列顺序为 [batch, height, width, channels]
比如 torch.channels_last
「NCHW」的排列顺序为 [batch, channels, height, width]
比如 torch.contiguous_format

通常来说,机器学习的不同 layer 的性能会受到输入 tensor 的内存排布的影响。比如对于 NVDIA 的 Tensor Core GPU 上的卷积层来说,输入 tensor 是 NHWC 的排布时计算效率最高 ¹。因为每个像素点的所有 channel 值同时被访问,NHWC 的数据排布更符合这种访问规律。但是,对于 Batchnorm 这样的 layer 来说,通常你会归纳 H、W 两个维度,则 NCHW 的排布计算效率是最高的。因此 PyTorch 提供了不同的 memory_format 以及它们之间灵活的转换,以方便用户能够在不同场景下使用最优的内存排布。

input = torch.randn(1, 10, 32, 32)
model = torch.nn.Conv2d(10, 20, 1, 1)
model = model.to(memory_format=torch.channels_last)
output = model(input)

相应地,Taichi 默认的 ti.Layout.AOS(array of struct)排布将 field 的 vector 元素中的每个标量在内存中连续排布。比如当用户想要同时对位置向量的 x, y, z 值加 1 的时候,AOS 的排布可以保证读取内存的效率是最高的。而且 Taichi 还通过 structure node 支持了 SOA(structure of arrays)排布和更高级(甚至是稀疏)的分层数据结构。比如下面这个例子展示了如何优化一个简单的应用场景。

如果我们已知该应用将逐个遍历很多大小为 8x8 的图片,我们就可以通过将 8x8 像素作为一个 cell 连续排布,从而最大化内存读写效率。同时由于 taichi 实现了算法和数据的分离,用户可以灵活地改变数据排布但无需改动 ti.kernel 内的计算代码,使得通过实验对不同的应用找到对其性能友好的数据排布变得非常简单。由于篇幅有限,感兴趣的读者可以通过下方链接的文档了解更多细节。taichi https://taichi-lang.cn/
郑重声明:文章内容来自互联网,纯属作者个人观点,仅供参考,并不代表本站立场 ,版权归原作者所有!

上一篇:学会这 6 个技巧,分分钟搭建炫酷场景 | Taichi 进阶秘籍

下一篇:如何有效驱逐和避免老鼠?

相关推荐

返回顶部