treevalue——Master Nested Data Like Tensor

首先，请和我一起高呼——“treevalue——通用树形结构建模工具 + 极简树形结构编程模型”。

咳咳，好久没更新了，这一次是真的好久不见，甚是想念。在之前的三期中，关于 treevalue 的核心特性等内容已经基本完成了讲述。因此本篇作为该系列的终章，将尝试用更高一层的视角来分析 treevalue ，以求进一步了解其核心思想与应用模式，并通过干货数据与实例展示来展现其真实能力。

闲话少叙，让我们开始吧！如果还没了解过 treevalue 的小伙伴们可以先去读一下之前的几篇文章：

• Treevalue(0x01)——功能概述
• Treevalue(0x02)——函数树化详细解析（上篇）
• Treevalue(0x03)——函数树化详细解析（下篇）

整体设计

概念与架构

想要完整地了解 treevalue ，首先还是需要从整体架构上来看看，如下图所示。

（treevalue的整体架构）

显而易见的一点， treevalue 是基于Python构建的，并且处于运算性能上的考虑，因此也大量采用了Cython来实现，该工具可以用一种介于Python和C/C++之间的语法进行编码，并以C/C++的形式编译为静态库，通过绕开一系列不必要的动态机制等方式来实现明显的加速。

在 treevalue 中，最底层的为 TreeStorage ，为数据层，主要对树状数据结构进行管理，并对上层提供最基本的接口。建立在之上的 TreeValue 为最关键的一个类，基于数据层进行了基本的封装，能实现 treevalue 的基本特性且具备进一步扩展的能力。树化（treelize）在之前的几篇文章中有过介绍，其作用在于将现有库的函数和类进行树化的扩展，使 treevalue 的特性可以被快速应用至现有工具上。工具部分（utilities）即为基于 TreeValue 构建的一些简单工具，支持了基本的树操作、函数式运算等功能。另一个极为重要的部分即为 FastTreeValue 类，其为 TreeValue 的子类，包含了大部分的运算符，并可以实现“装载即用”，只需要将现有的对象装入FastTreeValue即可批量访问其属性、批量调用其方法，因此 FastTreeValue 被最为广泛的使用。这里需要注意的是， treevalue 并未针对任何一个特定的现有库进行特殊化设计，而是通过泛用型设计，让诸如PyTorch、Numpy、Tensorflow在内的几乎全部接口均可以被快速扩展到树运算上。

总的来说，基于 treevalue ，可以将现有轮子中的函数与类扩展为支持树状运算的新形态，并可以让编程者基于这一扩展进行更加方便快捷的代码构建。

设计定位

在前面几篇该系列的文章里，我们都加上了“强化学习（Reinforcement Learning）”的标签，也是因此，此文才会被推送到关注了这一话题的用户手里。可是你们大概已经发现，无论是之前的三篇文章，还是我们的源码仓库，都并没有深度强化学习相关的内容，甚至在源代码中看不到 numpy 或 torch 等常见AI相关库的存在。想必各位读者也从之前的文章发出后便早已疑惑多时了。因此在这里我们来回答一个重要的问题—— treevalue的设计定位是什么？

其实仔细琢磨的话，答案也已经很明显了—— treevalue从设计之初就是一款具备充分泛用型，且易用性为主运行性能为辅的树数据结构计算框架。之前各个文章乃至代码中未大量涉及具体的深度强化学习内容也正是这个原因。在Treevalue(0x02)——函数树化详细解析（上篇）中介绍的 func_treelize 特性也是可以针对任意函数进行使用的，因此即便不是 torch 这样的经典的库，换上其他的库与其他的类、方法，也同样可以使用treevalue，且使用体验与原本的库基本一致，极为友好。甚至于深度强化学习以外的任意领域，只要有基于树结构的运算的地方， treevalue 就有它的用武之地——这也正是它设计上的核心理念，与真正的优势所在。

与同类产品的对比

基于上述的介绍，读者们可能会担心一个易用性、泛用性为主的库，会不会在其他上，尤其是在深度强化学习程序的特性支持与运行性能上存在劣势。这样的顾虑其实很正常，因为在大部分情况下，软件的泛用性和性能都是不容易做到兼顾的。但是 treevalue 是否存在这样的问题呢？因此本节将会通过与同类或相近产品的对比分析，来对这一问题进行回答。

同类或相近产品介绍

在 treevalue 开发与测试的阶段，我们发现其实在现有的其他库中，也有类似的树状运算支持。其中比较具有典型意义的有以下四个：

• dm-tree，DeepMind团队开发的轻量级树状运算库。
• Tianshou-Batch，清华大学机器学习组开发，基于PyTorch的深度强化学习库，其中Batch为重点设计的一个关键组件。
• jax-libtree，谷歌团队开发的机器学习库，其中libtree为内部的一个轻量级树运算封装库。
• torchbeast-nest，Meta研究院（原facebook研究院）开发，基于PyTorch的深度强化学习框架，其中nest为内部的一个轻量级树运算封装库。

特性对比

基于上面提到的四款同类产品及 treevalue ，所开源的代码及提供的文档，可以得出以下的分析。

具体来说：

• dm-tree为基于C++实现的独立轻量级库，不针对特定的框架，且支持简单的函数式运算。
• tianshou-Batch为中轻量级库，自建了一套针对PyTorch的树数据结构，实现了基本的增删查改功能，并且实现了包括stack、split在内的少量运算操作。
• jax-libtree为Google Jax库中的子模块，不针对特定框架，支持简单的函数式运算，还支持了诸如transpose这样的树结构运算。
• torchbeast-nest为Meta torchbeast中的树结构子模块，不针对特定框架，支持简单的函数式运算，此外还支持简单的二元函数扩展运算能力。 而相比之下， treevalue 为一款泛用型中量级库，且支持主要的函数式运算。其实现原理为自建树状数据结构，并基于这一结构展开运算。不仅如此，还支持了比jax-libtree更加丰富的结构运算（subside与支持自动检测结构的rise），也支持了比torchbeast-nest更灵活的函数扩展能力（func_treelize）。而 treevalue 最为强大的地方体现在treetensor中，只需要对部分 torch.Tensor 的方法进行特别支持后，剩下的全部方法均可在现有框架上实现，并保持和原有API一样的使用方式。

这一点，意味着对于基于 treevalue 的开发者而言，不再需要大规模逐个进行封装迁移，只需要针对个别较特殊的API进行特别实现，其他的可以直接批量快速映射，直接做到无一遗漏。实际上，对于兼容库的开发而言，往往实现了90%，甚至与99%的接口，都可能对实际使用带来较多的限制，而基于treevalue开发的应用而言，则在原理上便不存在这一问题，接口的扩展可以一步到位。而对于使用者而言，兼容库的使用体验将和原有库高度一致，原有的运算性质不会改变，且进一步支持了树状运算，而一步到位的兼容更是让使用者不会被不全的API所限制，使用体验极佳，甚至于可以将现有代码以极小规模的修改便实现从原有库向兼容库的迁移，很容易确保正常运行，并可以进一步简化代码实现。

性能对比

上文中提到， treevalue 是其中唯一的一款中量级库，这意味着从实现的角度来看，其需要为了支持高度的泛用性和易用性，而考虑更多的情况，设立更多的处理机制。因此，理论上 treevalue 会在性能上存在一定的劣势。然而事实真的如理论所言吗？让我们来看看接下来的性能分析数据。

在本次性能对比中，我们将根据同类产品的特点，将实验分为两组：

• 数据结构组——tianshou Batch和treevalue，比较对象为split（将单个Tensor进行拆分）、stack（将多个Tensor进行拼接）运算的性能
• 泛用运算组——dm-tree、jax-libtree、torchbeast-nest和treevalue，比较对象为flatten（将树结构展开为可逆的列表结构）、mapping（简单的函数式映射运算）运算的性能 数据结构组的性能测试结果如下所示，不难发现 treevalue 在张量运算的性能上具有明显优势。

（treevalue和tianshou Batch在torch.split操作上的性能对比，treevalue具有明显的优势）

（treevalue和tianshou Batch在torch.stack操作上的性能对比，treevalue具有明显的优势）

泛用运算组的性能测试结果如下图所示，其中dm-tree为其中最轻量级的库，却意外地拥有最低的性能；而jax-libtree和torchbeast-nest相比之下拥有不错的性能。但是，作为中量级库且包含大量易用性设计的treevalue，仍然在各种规模的数据上拥有性能优势——在小规模数据上优势极为明显，即便在大规模数据上依然可以较之jax-libtree保持微弱优势，且较之torchbeast-nest保持大比例的优势。

（treevalue、dm-tree、jax-libtree、torchbeast-nest在flatten运算上的性能对比，treevalue整体处于优势）

（treevalue、dm-tree、jax-libtree、torchbeast-nest在mapping运算上的性能对比，treevalue整体处于优势）

不仅如此，我们还针对其他的各类基础操作于运算进行了性能测试，更多的测试结果信息参见项目README.md。

经过一系列的测试，结果表明性能问题并未出现，甚至于相较于同类产品仍有性能优势。这一点对于中量级设计的库而言是极为难得的，这意味着无论从性能还是从使用体验来说，treevalue均具备全方位的优势。

案例对比

在上文中，我们在与同类产品的的对比中，展现了 treevalue 强大的泛用性与易用性，以及堪称优秀的运算性能。而实际上要想更具说服力地展现这一点，还是需要一些更加具体的例子。

首先是官方文档中的两组例子，分别用于Numpy和Scikit-Learn：

• Apply Into Numpy，这组例子展现了在实现同样涉及树数据结构的情况下，使用FastTreeValue的代码极为明显地短于原有实现，且简洁清晰程度远超以往。
• Apply Into Scikit-Learn，该组例子展现了基于原有的接口，通过极简的装饰扩展即可将PCA运算应用到整棵树上，得出整棵树的解，全程高度精简易懂。 此外，我们还准备了一个 treevalue 用于具体DRL项目的例子，这次我们选择了在深度强化学习领域最复杂的项目——AlphaStar中的相关数据处理函数 collate_fn ，这个函数的作用是在每个训练iteration之前对数据进行堆叠、填充和预处理。原始的代码充斥着大量的for循环和if-else分支控制，而通过使用 treevalue ，我们可以将上述所有操作完全用并行操作API重写，具体来说，两个版本的代码度量分析如下表所示：

这表明， treevalue 在复杂项目上的优势将更加显著，具体表现为代码行数缩减到1/3，逻辑复杂度大幅降低，可维护性明显提高——以及由此带来的编码用时大幅缩减。此外，在性能测试中，treevalue版代码相较于原始代码也并无劣势。至此， treevalue 的各项能力已经得到了充分的验证。

展望

treevalue 的各项能力，在上文中已经进行了充分的论述。因此，本节将采用开放式问题的方式，抛出一系列问题，也作为展望。欢迎读者一同参与讨论，以及如有更进一步的脑洞，也欢迎留言~~

问题1：你认为， treevalue 仅仅只是一个库？一组操作工具？还是一整套运算模型？

问题2： treevalue运算模型的核心是什么？此处核心指的是，基于该核心可以直接或间接衍生出几乎全部现有特性。

问题3：在之前的文章中，已经详细描述了函数的树化扩展（treelize）机制，效果为将普通函数扩展为支持树运算的函数。除此之外，是否还有其他类似的可扩展对象？是否有可以基于函数树化的上位可扩展对象？

问题4：如果需要对现有的数据模型类（例如torch.Tensor、np.ndarray等）进行树化扩展，需要注意哪些问题？工厂类（工厂模式构建的类）呢？工具类（集成静态工具的类，多见于Java）呢？对模块（module）的扩展呢？上述的各种扩展机制是否可以归纳为另一种统一的运算模型？

问题5：treevalue相较于同类产品的性能优势虽然普遍存在，但依然没有做到全方位拉开差距。有哪些可能的优化点？其中哪些是可以针对具体应用情况作针对性优化的？

问题6：如何通过扩展让treevalue支持任意类型的数据容器（包括list、tuple，甚至自建数据模型）？眼下的主要障碍在哪里？是在技术上还是运算模型上？

问题7：你还能想到其他的 treevalue + 机器学习的应用吗？图神经网络数据结构的表示是否可以实现？