DI-treetensor

PyPI - Python版本代码行数

GitHub提交活动

treetensor是一个通用的基于树的张量结构，主要由OpenDILab贡献者开发。

几乎所有操作都可以以树的形式方便地支持，以简化计算基于树结构时的处理过程。

安装

你可以使用pip命令行从官方PyPI网站简单安装它。

pip install di-treetensor

有关安装的更多信息，你可以参考安装。

文档

详细文档托管在https://opendilab.github.io/DI-treetensor。

目前只提供英文版本，中文文档仍在开发中。

快速开始

你可以基于FastTreeValue轻松创建树值对象。

import builtins
import os
from functools import partial

import treetensor.torch as torch

print = partial(builtins.print, sep=os.linesep)

if __name__ == '__main__':
    # 创建树张量
    t = torch.randn({'a': (2, 3), 'b': {'x': (3, 4)}})
    print(t)
    print(torch.randn(4, 5))  # 创建普通张量
    print()

    # 树的结构
    print('树的结构')
    print('t.a:', t.a)  # t.a是原生张量
    print('t.b:', t.b)  # t.b是树张量
    print('t.b.x', t.b.x)  # t.b.x是原生张量
    print()

    # 数学计算
    print('数学计算')
    print('t ** 2:', t ** 2)
    print('torch.sin(t).cos()', torch.sin(t).cos())
    print()

    # 反向计算
    print('反向计算')
    t.requires_grad_(True)
    t.std().arctan().backward()
    print('t的梯度:', t.grad)
    print()

    # 原生操作
    # 所有操作都可以像原始`torch`一样使用
    print('原生操作')
    print('torch.sin(t.a)', torch.sin(t.a))  # 原生张量的sin

结果应该是

<Tensor 0x7f0dae602760>
├── a --> tensor([[-1.2672, -1.5817, -0.3141],
│                 [ 1.8107, -0.1023,  0.0940]])
└── b --> <Tensor 0x7f0dae602820>
    └── x --> tensor([[ 1.2224, -0.3445, -0.9980, -0.4085],
                      [ 1.5956,  0.8825, -0.5702, -0.2247],
                      [ 0.9235,  0.4538,  0.8775, -0.2642]])

tensor([[-0.9559,  0.7684,  0.2682, -0.6419,  0.8637],
        [ 0.9526,  0.2927, -0.0591,  1.2804, -0.2455],
        [ 0.4699, -0.9998,  0.6324, -0.6885,  1.1488],
        [ 0.8920,  0.4401, -0.7785,  0.5931,  0.0435]])

树的结构 t.a: tensor([[-1.2672, -1.5817, -0.3141], [ 1.8107, -0.1023, 0.0940]]) t.b: <Tensor 0x7f0dae602820> └── x --> tensor([[ 1.2224, -0.3445, -0.9980, -0.4085], [ 1.5956, 0.8825, -0.5702, -0.2247], [ 0.9235, 0.4538, 0.8775, -0.2642]])

t.b.x tensor([[ 1.2224, -0.3445, -0.9980, -0.4085], [ 1.5956, 0.8825, -0.5702, -0.2247], [ 0.9235, 0.4538, 0.8775, -0.2642]])

数学计算 t ** 2: <Tensor 0x7f0dae602eb0> ├── a --> tensor([[1.6057, 2.5018, 0.0986], │ [3.2786, 0.0105, 0.0088]]) └── b --> <Tensor 0x7f0dae60c040> └── x --> tensor([[1.4943, 0.1187, 0.9960, 0.1669], [2.5458, 0.7789, 0.3252, 0.0505], [0.8528, 0.2059, 0.7699, 0.0698]])

torch.sin(t).cos() <Tensor 0x7f0dae621910> ├── a --> tensor([[0.5782, 0.5404, 0.9527], │ [0.5642, 0.9948, 0.9956]]) └── b --> <Tensor 0x7f0dae6216a0> └── x --> tensor([[0.5898, 0.9435, 0.6672, 0.9221], [0.5406, 0.7163, 0.8578, 0.9753], [0.6983, 0.9054, 0.7185, 0.9661]])

反向计算 t的梯度: <Tensor 0x7f0dae60c400> ├── a --> tensor([[-0.0435, -0.0535, -0.0131], │ [ 0.0545, -0.0064, -0.0002]]) └── b --> <Tensor 0x7f0dae60cbe0> └── x --> tensor([[ 0.0357, -0.0141, -0.0349, -0.0162], [ 0.0476, 0.0249, -0.0213, -0.0103], [ 0.0262, 0.0113, 0.0248, -0.0116]])

原生操作 torch.sin(t.a) tensor([[-0.9543, -0.9999, -0.3089], [ 0.9714, -0.1021, 0.0939]], grad_fn=<SinBackward>)


若需更多快速入门说明和进一步用法，请查看：

* [快速入门](https://opendilab.github.io/DI-treetensor/main/tutorials/quick_start/index.html)

## 扩展

如果您需要将`treevalue`对象转换为可运行的源代码，可以使用[potc-treevalue](https://github.com/potc-dev/potc-treevalue)插件，安装命令如下

pip install DI-treetensor[potc]


在potc中，您可以将对象转换为可运行的Python源代码，之后可以通过Python解释器加载为对象，如下图所示

![potc系统](https://yellow-cdn.veclightyear.com/835a84d5/98e9c529-9a09-4023-bf6f-e4ed1e92d13e.svg)

更多信息，您可以参考

- [potc-dev/potc](https://github.com/potc-dev/potc)
- [potc-dev/potc-treevalue](https://github.com/potc-dev/potc-treevalue)
- [potc-dev/potc-torch](https://github.com/potc-dev/potc-torch)
- [Potc插件安装](https://opendilab.github.io/DI-treetensor/main/tutorials/plugins/index.html#potc-support)

## 贡献

我们感谢所有为改进DI-treetensor而做出的贡献，无论是逻辑还是系统设计。更多指南请参考CONTRIBUTING.md。

用户可以加入我们的[slack交流频道](https://join.slack.com/t/opendilab/shared_invite/zt-v9tmv4fp-nUBAQEH1_Kuyu_q4plBssQ)，或联系核心开发者[HansBug](https://github.com/HansBug)进行更详细的讨论。

## 许可证

`DI-treetensor`在Apache 2.0许可下发布。