TreeValue——一个通用树状数据结构与函数计算库
Treevalue v1.0.0版本已经于2021年10月24日正式发布，欢迎下载体验：opendilab / treevalue。
这算是treevalue的第一个正式实用化版本，本文将会对其主要功能和特性进行一个概述。
一个直观地展示
设想这样一个实际的应用场景，我们需要使用numpy对机器学习中的一批样本进行预处理，并组装成一个训练用的mini-batch。一个数据样本的格式如下面的代码所示，即函数 get_data 的返回值：
import numpy as np
T, B = 3, 4
def get_data():
return {
'a': np.random.random(size=(T, 8)),
'b': np.random.random(size=(6,)),
'c': {
'd': np.random.randint(0, 10, size=(1,))
}
}
如果使用最常见的写法，大概会是这样的
# without_treevalue.py
import numpy as np
T, B = 3, 4
def without_treevalue(batch_):
mean_b_list = []
even_index_a_list = []
for i in range(len(batch_)):
for k, v in batch_[i].items():
if k == 'a':
v = v.astype(np.float32)
even_index_a_list.append(v[::2])
elif k == 'b':
v = v.astype(np.float32)
transformed_v = np.power(v, 2) + 1.0
mean_b_list.append(transformed_v.mean())
elif k == 'c':
for k1, v1 in v.items():
if k1 == 'd':
v1 = v1.astype(np.float32)
else:
print('ignore keys: {}'.format(k1))
else:
print('ignore keys: {}'.format(k))
for i in range(len(batch_)):
for k in batch_[i].keys():
if k == 'd':
batch_[i][k]['noise'] = np.random.random(size=(3, 4, 5))
mean_b = sum(mean_b_list) / len(mean_b_list)
even_index_a = np.stack(even_index_a_list, axis=0)
return batch_, mean_b, even_index_a
而当我们有了treevalue库之后，完全一致的功能可以被这样简短的代码实现
# with_treevalue.py
import numpy as np
from treevalue import FastTreeValue
T, B = 3, 4
power = FastTreeValue.func()(np.power)
stack = FastTreeValue.func(subside=True)(np.stack)
split = FastTreeValue.func(rise=True)(np.split)
def with_treevalue(batch_):
batch_ = [FastTreeValue(b) for b in batch_]
batch_ = stack(batch_)
batch_ = batch_.astype(np.float32)
batch_.b = power(batch_.b, 2) + 1.0
batch_.c.noise = np.random.random(size=(B, 3, 4, 5))
mean_b = batch_.b.mean()
even_index_a = batch_.a[:, ::2]
batch_ = split(batch_, indices_or_sections=B, axis=0)
return batch_, mean_b, even_index_a
可以看到，实现一段同样的基于树结构的业务逻辑，有了treevalue库的辅助后代码变得极为简短和清晰，也更加易于维护。
这正是treevalue的最大亮点所在，接下来的章节中将会对其主要功能和特性进行概述，以便读者对这个库有一个整体的了解和认识。
树结构及其基本操作
在treevalue库中，我们提供一种核心数据结构—— TreeValue 类。该类为整个treevalue的核心特性，后续的一系列操作都是围绕 TreeValue 类所展开的。
首先是 TreeValue 对象的构造（使用的是增强版子类 FastTreeValue ，关于 TreeValue 类与 FastTreeValue 类的区别可以参考文档，本文不作展开），只需要将dict格式的数据作为唯一的构造函数参数传入即可完成TreeValue的构造
from treevalue import FastTreeValue
t = FastTreeValue({'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}})
# <FastTreeValue 0x7f135a5ada30>
# ├── a --> 1
# ├── b --> 2
# └── x --> <FastTreeValue 0x7f135a5ad970>
# ├── c --> 3
# └── d --> 4
不仅如此， TreeValue 类还提供了树状结构的几种基本操作接口供使用
from treevalue import FastTreeValue
t = FastTreeValue({'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}})
## get value/node
t.a # 1
t.x # <FastTreeValue 0x7f135a5ad970>
# ├── c --> 3
# └── d --> 4
t.x.d # 4
## set value/node
t.x.d = 35
# t after t.x.d = 35
# <FastTreeValue 0x7f135a5ada30>
# ├── a --> 1
# ├── b --> 2
# └── x --> <FastTreeValue 0x7f135a5ad970>
# ├── c --> 3
# └── d --> 35
## delete value/node
del t.b
# t after del t.b
# <FastTreeValue 0x7f135a5ada30>
# ├── a --> 1
# └── x --> <FastTreeValue 0x7f135a5ad970>
# ├── c --> 3
# └── d --> 35
## contains key or not
'a' in t # True
'd' in t # False
'd' in t.x # True
## size of node
len(t) # 2, 'a' and 'x'
len(t.x) # 2, 'c' and 'd'
## iterate node
for k, v in t.x:
print(k, '-->', v)
# c --> 3
# d --> 35
以上是 TreeValue 类的几种常见基本操作，支持了最基本的增删查改等。
树的可视化
当一个 TreeValue 对象被构造出来后，我们如何比较直观地去观察其整体结构呢？有两种方式可以对 TreeValue 进行可视化：
- 通过
print进行快速文字打印 - 通过
treevalue graph命令行工具进行图像导出
实际上对于第一种情况，在上一节中已经有了展示，在此展示一个更加复杂的案例
# test_simple.py
import torch
from treevalue import FastTreeValue
t = FastTreeValue({
'a': torch.randn(2, 3), # objects with multiple lines
'b': torch.randn(3, 1, 2),
'x': {
'c': torch.randn(3, 4),
}
})
t.x.d = t.x # nested pointer
print(t)
输出结果如下，可以看到诸如 torch.Tensor 这样多行的对象也一样可以被有序地排版输出，且对于存在嵌套引用的情况，输出时也可以被准确地识别出来，避免无限循环打印
<FastTreeValue 0x7f642057bd00>
├── a --> tensor([[ 0.1050, -1.5681, -0.2849],
│ [-0.9415, 0.2376, 0.7509]])
├── b --> tensor([[[ 0.6496, -1.3547]],
│
│ [[ 1.2426, -0.2171]],
│
│ [[-0.7168, -1.4415]]])
└── x --> <FastTreeValue 0x7f642057bd30>
├── c --> tensor([[-0.6518, 0.4035, 1.0721, -0.6657],
│ [ 0.0252, 0.4356, 0.1284, -0.3290],
│ [-0.6725, 0.2923, 0.0048, 0.0704]])
└── d --> <FastTreeValue 0x7f642057bd30>
(The same address as <root>.x)
除了基于文本的可视化外，我们还提供了命令行工具以进行图像导出。例如上面的代码，我们可以用如下的命令行导出图像
treevalue graph -t 'test_simple.t' -o 'test_graph.png'
此外，对于更复杂的情况，例如这样的一份源代码
# test_main.py
import numpy as np
from treevalue import FastTreeValue
tree_0 = FastTreeValue({
'a': [4, 3, 2, 1],
'b': np.array([[5, 6], [7, 8]]),
'x': {
'c': np.array([[5, 7], [8, 6]]),
'd': {'a', 'b', 'c'},
'e': np.array([[1, 2], [3, 4]])
},
})
tree_1 = FastTreeValue({
'aa': tree_0.a,
'bb': np.array([[5, 6], [7, 8]]),
'xx': {
'cc': tree_0.x.c,
'dd': tree_0.x.d,
'ee': np.array([[1, 2], [3, 4]])
},
})
tree_2 = FastTreeValue({'a': tree_0, 'b': tree_1, 'c': [1, 2], 'd': tree_1.xx})
可以通过以下的命令行导出为图像，不难发现对于共用节点和共用对象的情况也都进行了准确地体现（如需进一步了解，可以执行 treevalue graph -h 查看帮助信息）
treevalue graph -t 'test_main.tree_*' -o 'test_graph.png' -d numpy.ndarray -d list -d dict
以上是对于 TreeValue 对象的两种可视化方法。
函数的树化
在treevalue中，我们可以快速地将函数进行装饰，使之可以支持 TreeValue 对象作为参数进行运算。例如下面的例子
# test_gcd.py
from treevalue import FastTreeValue, func_treelize
@func_treelize(return_type=FastTreeValue)
def gcd(a, b): # GCD calculation
while True:
r = a % b
a, b = b, r
if r == 0:
break
return a
if __name__ == '__main__':
t1 = FastTreeValue({'a': 2, 'b': 30, 'x': {'c': 4, 'd': 9}})
t2 = FastTreeValue({'a': 4, 'b': 48, 'x': {'c': 6, 'd': 54}})
print("Result of gcd(12, 9):", gcd(12, 9))
print("Result of gcd(t1, t2):")
print(gcd(t1, t2))
将整数之间的最大公因数进行了装饰后，可以形成兼容 TreeValue 对象的函数，且不会影响普通对象的运算结果，因此上述代码的输出结果如下所示
Result of gcd(12, 9): 3
Result of gcd(t1, t2):
<FastTreeValue 0x7f53fa67ff10>
├── a --> 2
├── b --> 6
└── x --> <FastTreeValue 0x7f53fa67ff40>
├── c --> 2
└── d --> 9
树结构的运算
除了 TreeValue 自带的一系列基本操作之外，treevalue库还提供了一些常用的树结构运算函数。例如如下的四种：
- map——值映射运算
- reduce——值归纳运算
- subside——顶层结构下沉运算
- rise——值结构提取上浮运算
值映射运算（map）
TreeValue 对象的值映射运算和列表类型的map运算类似，会产生一个同结构且值为映射值的新 TreeValue 对象，例如下面的案例
from treevalue import FastTreeValue
t1 = FastTreeValue({'a': 2, 'b': 30, 'x': {'c': 4, 'd': 9}})
t2 = t1.map(lambda x: x * 2 + 1)
t1 和 t2 的图像如下
值归纳运算（reduce）
TreeValue 对象的值归纳运算和 functools.reduce 函数的功能类似，可以将树结构以子树为单位进行归纳，最终计算出一个结果来，例如下面的案例
from treevalue import FastTreeValue
t1 = FastTreeValue({'a': 2, 'b': 30, 'x': {'c': 4, 'd': 9}})
# sum of all the values in t1
t1.reduce(lambda **kws: sum(kws.values())) # 45
可以快捷的实现整棵树的求和运算。
顶层结构下沉运算（subside）
TreeValue 还可以支持将其上层结构进行解析后，下沉到节点值上，形成一棵新的树，例如下面的案例
from treevalue import FastTreeValue
dt = {
'i1': FastTreeValue({'a': 1, 'b': 2, 'x': {'c': 3}}),
'i2': (
FastTreeValue({'a': 7, 'b': 4, 'x': {'c': 6}}),
FastTreeValue({'a': 11, 'b': 9, 'x': {'c': -3}}),
),
}
t = FastTreeValue.subside(dt)
下沉后产生的树 t 的图像为
值结构提取上浮运算（rise）
上浮运算（rise）为下沉运算的逆运算，可以从值节点中提取共同结构至 TreeValue 对象外，例如如下的案例
from treevalue import FastTreeValue, raw
t = FastTreeValue({
'a': raw({'i1': 1, 'i2': (7, 11)}),
'b': raw({'i1': 2, 'i2': (4, 9)}),
'x': {
'c': raw({'i1': 3, 'i2': (6, -3)}),
}
})
dt = t.rise()
dt_i1 = dt['i1']
dt_i2_0, dt_i2_1 = dt['i2']
对象 dt 是一个字典类型的对象，包含 i1 和 i2 两个键，其中 i1 为一个 TreeValue 对象， i2 为一个长度为2，由 TreeValue 对象构成的元组。因此 dt_i1 、 dt_i2_0 和 dt_i2_1 的图像为
后续预告
本文作为一个对于treevalue现有主要功能的一个概述，受限于篇幅暂时做不到深入全面的讲解内部原理和机制。因此后续会考虑继续出包括但不限于以下的内容：
- treevalue的整体体系结构
func_treelize树化函数的原理与机制- treevalue的一些神奇的用法与黑科技
- treevalue的优化历程与经验
- treevalue的具体实战案例
敬请期待，同时欢迎了解其他OpenDILab的开源项目：https://github.com/opendilab。