Treevalue(0x01)——功能概述

HansBug

发布于 2021-11-04 16:56:22

2570

发布于 2021-11-04 16:56:22

文章被收录于专栏：HansBug's Lab

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 的图像为