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。

巴特西

Treevalue(0x01)——功能概述