Skip to content

SeetaResearch / Dragon

Go to a project
Developer guide

Developer guide

Last edited by Ting PAN
Page history

1. 计算核心

1.1 技术指南

  • 主要使用C++ (C++11或更高特性), 大量使用宏和泛型的两者组合, 以减少重复代码

  • CPU计算加速优化:

    • 优先使用Eigen(third_party/eigen), 其底层抽象了X86和ARM64指令集
    • 对于X86, 然后考虑OpenMP，看情况可以手写SSE, AVX2, FMA, AVX512
    • 对于ARM64, 优先适配NPU的SDK, 最后考虑手写NEON
    • 根据时间节点，考虑 MKL / MKLDNN 的集成，CPU算子前期主要是配合单元测试

  • CUDA计算加速优化:

    • 尽可能为每个算子实现Kernel, 对于复杂的并行算法(广播、规约、排序),

    可以使用CUB(third_party/cub)和cuBLAS/Thrust(CUDA集成)

    • 尽可能利用每一个cuDNN接口, 编写cuDNN加速版可选算子
    • 尽可能利用GPU的Cache/SharedMem，减少临时空间的使用
    • 尽可能通过预定义使用新GPU架构的Intrinsic指令, 比如sm35以上的__ldg,

    sm53 以上的half相关指令，以及安培架构新增的bfloat16相关指令

    • 尽可能不要强制同步当前执行的CUDA Stream

  • 分布式计算优化:

    • 集群通信基础库由third_party/mpi负责, 我们默认用的是OpenMPI
    • 用NCCL实现支持的Collective MPI操作, 实现GPUDirect通过RDMA直接加速
    • 由于pip预编译版本禁用了OpenMPI的RDMA特性, 暂不考虑MPI对Infiniband优化

  • 静态计算图

    • 输入为满足SSA形式的proto结构, 我们称为GraphIR, 包含上下文的算子和一些执行信息
    • 目前实现了三个级别的优化, 分别是:
    • 图剪枝 (O1, 移除求解最终输出变量不需要的Op或者Op的部分输出)
    • 追加: 自适应的In-place (O2, 输出直接覆盖只被引用一次的输入)
    • 追加: 模拟垃圾回收(GC)机制 (O3, 用FIFO/LIFO的Buffer池改写全部中间输出)

    默认优化级别为O3, 模拟GC在没有梯度计算时等价于PyTorch的no_grad执行时的内存优化,

    需要计算梯度时, 前向计算回滚到O2优化, 而反向计算的O3优化则实现了主流框架的梯度内存共享特性

  • 动态计算图

    • 前向计算输入为Op的proto结构, 反向计算输入为全部前向Op
    • 仅反向部分可用静态计算图形式构建, 由于O2优化此时由用户定义, 不满足SSA,

    故O1的剪枝优化只能移除无关的叶子变量, 无法分析中间变量, O3优化同静态计算图

    • Tensor垃圾回收目前利用PythonGC实现, 在计算核心中的管理方式同静态计算图

  • 内存管理

    • UnifiedMemory统一管理Host和Device内存,  并由状态机实现data间的同步
    • 目前不支持non-contigous 内存布局
    • 一般由Tensor持有自己的UnifiedMemory, 或映射到其它的Tensor的UnifiedMemory

  • Workspace

    • 维护Tensor, Graph, 以及因动态图而cache参数的Op
    • 等价于tf.Graph, Caffe2的Workspace, 但我们沿用了Caffe2的命名
    • 用于可命名资源的隔离管理及上下文视图, 实现DL框架的MVC模式

  • 算子

    • 我们的Op实现为 抽象计算设备/数据类型 的计算 Dispatcher
    • 主要任务是处理静态/动态参数、形状推理,

    最后根据输入Tensor设备/数据类型调用对应OpKernel函数实现

    • OpKernel、MathKernel与Op独立实现，便于多个Op调用相同目标的计算

  • C++代码规范

    推荐编辑器安装clang-format插件, 每次保存时自动格式化

1.2 开发路线图

  • 高阶梯度 (PCL开源竞赛题)
    • 对列出的全部算子，实现二阶梯度计算
    • Div
    • Sin, Cos, Tanh, Relu, Log, Square, Reciprocal, Abs, Clip
    • Dropout, Sigmoid, Sqrt, Rsqrt, FullyConnected, Conv2d
    • 对列出的部分算子，实现无限阶梯度计算
    • Add, Sub, Mul
  • 重构静态计算图的优化
    • 参考ONNX和传统C/C++编译器, 细粒度化每一个PASS
    • 重写亚线性内存优化, 引入代价函数实现局部计算图的不同变换
  • 利用cuDNN 7.5引MultiHeadAttn 接口编写优化Transformer的Op
  • 考虑设计更好的动态计算变量池方案
    • 目前前向计算为FIFO, 反向为LIFO, 对于NAS任务完全失效
  • 增加TopKOp, 需要对数学库增加cuda排序(Radix/Heap)的实现
  • Mobile平台的模型推理
    • 实现部分内置算子的INT8计算 (NEON)

2. Python接口

2.1 技术指南

  • Pybind

    • 为Tensor 、Workspace 等C++资源类绑定PyObject,

    方便通过PythonGC为框架前端设计更灵活的接口

    • 操纵NumPy或DLPack的Python对象，实现ZeroCopy的内存映射
    • 在Python中使用C++实现的Protobuf Message, 以减少动态图传递、编辑IR的开销

  • 独立的编译模块

    • Pybind模块代码与框架主体分离在两个不同的动态库中,

    否则在Windows下, Python解释器加载的扩展DLL, 不能共享整合DLL内Registry

  • 静态计算 (AutoGraph)

    • 执行模式为GRAPH_MODE
    • dragon.Tensor中记录从叶子变量计算到本变量的全部OpDef,

    以及本变量对其它变量的偏导目标

    • 根据前向计算的OpDef和偏导目标生成反向传播的OpDef及

    参数更新的OpDef, 加上执行信息一起注入进GraphDef

    • 将GraphDef传递给计算核心在当前Workspace构建优化的Graph
    • 返回一个绑定Graph执行和Feed输入的回调函数给用户反复调用
    • 装饰器模式: 通过inspect模块获取装饰器标记的函数的输入参数信息,

    将输入替换为dragon.Tensor执行标记的函数, 取输出变量中的记录的OpDef

    此过程等价tf 2.0引入的tf.function、pytorch 1.0引入的torch.jit.script

  • 动态计算 (EagerExecution)

    • 执行模式为EAGER_MODE
    • dragon.GradientTape中记录上下文管理器进入至退出期间,

    所有watched/trainable的dragon.EagerTensor, 执行所使用的OpDef

    • 根据前向计算的OpDef和偏导目标生成反向传播的OpDef,

    并由当前Workspace逐个执行

    • 参数更新器对已计算出梯度的参数, 逐个执行Update算子
    • dragon.EagerTensor参与PythonGC的引用计数,

    引用归零且未被反向传播锁定则进入变量池

  • 内置IO

    • 层次化的设计
    • HDD/SSD->Bytes (DataReader)
    • Bytes->Image/Label (DataTransformer)
    • Image/Label -> ImageBlob/LabelBlob (DataIterator)
    • 三者位于不同进程中, 异步缓存数据, 通过同步队列传递数据
    • DataReader
    • 负责从硬盘进行分布式的非连续字节读取
    • 需要预先序列化数据, 默认的Message结构为我们设计的KPLRecord
    • 两种Shuffle策略: Example-wise和Chunk-wise

    前者为每个epoch后逐样本打乱, 后者则是逐块打乱, 而块内顺序读取

    Chunk-wise保证了HDD能够稳定支持单机8卡分布式ImageNet训练

    • DataTransformer
    • 数据解码、数据增强、标签处理
    • 不同DL任务需要重写该类, 默认提供图像分类的Transformer
    • DataIterator
    • 从DataTransformer中取出数据组Batch

  • 参数更新器

    • 每一个slot控制一个param group, 组内学习率, clip norm, weight decay共享
    • apply_gradients 在静态计算下, 收集参数梯度二元组, 反之直接执行计算
    • 分布式的ProcessGroup进入上下文管理器时, 增加AllReduceOp处理数据并行

  • 虚拟框架: TensorFlow

    • 不支持也无计划支持tf.compat.v1
    • tf 2.0大部分组件与dragon重合, 故仅仅是alias
    • tf.Graph -> dragon.Workspace
    • tf.GradientTape -> dragon.GradientTape
    • tf.Tensor -> dragon.Tensor
    • tf.EagerTensor -> dragon.EagerTensor
    • tf.function -> dragon.function
    • tf.gradients -> dragon.gradients
    • tf.ops -> dragon.ops
    • tf.Variable 目前确定为 tf.ResourceVariable,

    在任何计算模式下都是dragon.EagerTensor的子类,

    可直接可用于静态计算

    • tf.keras部分基本组件进行了重新设计
    • tf.keras.Network 和 tf.keras.Model对tf.keras.Layer

    的堆叠功能下沉到tf.keras.Layer, 与tf.Module设计理念一致

    • 以上三者的自动命名支持通过__setattr__时获取,

    与torch.nn.Module设计理念一致, 简化name参数的传递

    • tf.keras.optimizers.Optimizer在dragon.updaters.Updater基础上,

    增加了对Regularizer的处理, slot的基准weight decay取所有regularizers

    的最小值, 非最小值则使用UpdateOp的decay_mult参数实现

  • 虚拟框架: PyTorch

    • torch.Tensor独立于dragon.EagerTensor, 我们没有办法统一两者的接口,

    受此影响, torch.ops不可对dragon.ops做简单地alias, 需要重新封装一次

    • 动态计算接口针对PyTorch进行了部分微调
    • OpDef默认通过torch.Tensor存储和传递 (与静态计算类似)

    GradientTape 仅仅在torch.jit.trace接口中进行全局记录

    • zero_grad不是必须的, 计算核心并不会主动累积梯度, 也就不需要置零
    • 由于Python不支持函数重载, 对PyTorch部分从C++导出接口解耦:
    • torch.max -> torch.max, torch.argmax, torch.maximum
    • torch.min -> torch.min, torch.argmin, torch.minimum
    • 由于不支持non-contigous特性的, 与dim、stride相关的Op存在计算/存储代价
    • view, reshape, squeeze, unsqueeze
    • perimute
    • slice, expand

  • 算子插件

    • Op可从外部动态库中加载和注册, 无须重新编译框架
    • 通过C++/CUDA实现不支持的算子，提升特定算法平台的训练/推理性能

  • Python代码规范

    • 遵守PEP8规范
    • 每个Python代码文件, 必须经过仓库内规定的.flake8检查
    • 避免使用__all__, 如需分隔用户/开发者的接口视角, 使用_api和core文件夹结构

  • Python注释规范

    • 使用NumPy风格的编写docstring, 参数和返回值的数据类型必须注解精确
    • 普通注释首字母大写, 包含主语的句子需要以句号结尾
    • 语言表达需简洁扼要

2.2 开发路线图

  • 与PyTorch 1.5对齐接口
  • 与TensorFlow 2.2对齐接口
  • 与TensorLayer对齐接口
  • 完善TensorRT对INT8量化、动态形状接口
  • 持续跟踪DALI新的数据处理算子

3. 文档

3.1 技术指南

  • Sphinx-doc
    • 通过Sphinx解析docstring生成Python API文档
    • Sphinx主题具有高度可定制性, Tensorflow/PyTorch主页皆做了深度的设计
    • 我们设计了自己的主题: sphinx-seeta-theme，借鉴了TensorFlow主题中一些交互元素
  • reStructuredText
    • 使用rst格式编写控制自动的文档生成源文件
    • rst相比markdown，支持生成跨文件的目录结构树(TOC)

3.2 开发路线图

  • 开发Sphinx主题: 增加二级导航栏, 一级导航栏的二级菜单, 增加左侧边栏的可定制性
  • 补充Getting started 和 Tutorials

4. 测试

4.1 技术指南

  • TBD

4.2 开发路线图

  • 编写Op的单元测试, 考虑同时将所有的虚拟框架一起测试
  • 编写核心功能的单元测试

5. 算法平台开发

5.1 图像分类

  • 支持多种backbone和attention模块的组合
  • 支持多种损失函数，如Focal Loss、Dice Loss等
  • 优化数据/IO，支持分布式训练/分布式BatchNorm
  • 支持ONNX模型导出与TensorRT推理优化
  • 使用PyTorch或Tensorflow虚拟接口开发

5.2 目标检测

  • 复杂算子的实现: Deformable Conv
  • Backbone的补充: HRNet, ResNeSt
  • 检测算法的补充: Cascade R-CNN, FCOS, NAS-FPN
  • 使用PyTorch或Tensorflow虚拟接口开发

5.3 语义分割

  • 支持多种backbone和attention模块的组合
  • 支持多种损失函数，如Focal Loss、Dice Loss等
  • 优化数据/IO，支持分布式训练/分布式BatchNorm
  • 支持ONNX模型导出与TensorRT推理优化
  • 使用PyTorch或Tensorflow虚拟接口开发

5.4 GAN

  • 实现DCGAN/ProgressiveGAN/StyleGAN/BigGAN
  • 支持CIFAR10/ImageNet/CelebA/Celeb-HD等数据集
  • 使用PyTorch或Tensorflow虚拟接口开发

5.5 GCN

  • 构建常用的GCN模块
  • 处理Graph消息传递问题，并尝试优化执行效率
  • 在常见数据上完成benchmark，和其他GCN框架比较精度和速度
  • 使用PyTorch或Tensorflow虚拟接口开发