cstsunfu/dlk

A PyTorch Based Deep Learning Quick Develop Framework. One-Stop for train/predict/server/demo

Deep Learning toolKit (dlk)

Don't Repeat Yourself

简体中文 | English

dlk                                  --
├── adv_method                       -- adversarial training method like free_lb, fgm, etc.
├── callback                         -- callbacks, like checkpoint, early_stop, etc.
├── data                             -- data processor part
│   ├── data_collate                 -- data collate for collate a batch of data from dataset to dataloader
│   ├── datamodule                   -- the datamodule a.k.a lightning.LightningDataModule
│   ├── dataset                      -- the dataset inherit the torch.Dataset
│   ├── postprocessor                -- the tasks postprocessor
│   ├── processor                    -- the default processor, which scheduler the subprocessors
│   └── subprocessor                 -- the subprocessors like tokenizer, token2id, etc.
├── display                          -- the tasks display setting
├── imodel                           -- the integrated model, which a.k.a the lightning.LightningModule
├── initmethod                       -- the initmethod, some classic parameter init methods
├── nn                               -- builtin nn modules
│   ├── base_module.py               --
│   ├── layer                        --
│   │   ├── decoder                  --
│   │   ├── embedding                --
│   │   ├── encoder                  --
│   │   └── token_gen_decoder        --
│   ├── loss                         --
│   ├── model                        --
│   ├── module                       --
│   └── utils                        --
├── token_sample                     -- for text generate, different sample strategies
├── optimizer                        -- optimizers
├── scheduler                        -- learning rate schedulers
├── trainer                          -- the trainer, a.k.a lightning.Trainer
├── utils                            --
├── preprocess.py                    -- preprocess datas for train|predict|demo|etc.
├── train.py                         -- train entry
├── online.py                        --
├── predict.py                       -- just predict a bunch of data using the pretrained model
├── server.py                        -- deploy this to server your pretrained model
├── demo.py                          -- demo main
└── version.txt                      --

• Install
• Demo
  • Grid Search
  • Task Demo
• Usage and Feature
  • 使用方法
  • 模块注册
  • 部分内置模块介绍
    • callback
    • 虚拟对抗训练
    • 复杂训练控制
    • 文本生成
  • 实现你自己的模型
  • More Document

虽然最近的一年多通用大模型吸引了大部分人的注意力，但是相信很多人已经意识到任务导向的模型在现阶段仍有其不可替代的一面，而且这些模型在处理某些特定任务时具有更好的可靠性和更高的效率，特别是这些模型可以实现一些Agent来与LLM进行配合。

任务导向的模型开发实际上不像LLM一样可以“一招鲜吃遍天”，而是每个任务的模型都需要针对性的开发，而在工作中我们经常需要对深度神经网络模型进行快速实验，搜索最优结构和参数，并将最优模型进行部署，有时还需要做出demo进行验证.

首先是不同任务的开发实际上有很大一部分是重复的，而同一个任务的训练、预测、部署和demo这几个步骤的核心代码也是一致的，但是在实现上都需要一定的的改动，如果每个步骤都独立开发的话，会使得整个过程非常割裂，而这造成的代码冗余对于长期的代码维护是灾难性的。

DLK是一个使用lightning的Trainer，intc为config管理系统的集模型训练、参数（架构）搜索、模型预测、模型部署和demo为一身，对于同一个模型实现这些功能只需要依赖一份代码，大大降低开发和维护成本.

同时DLK作为一个通用的训练框架，我们的各种训练技巧和增强方法也可以非常方便的用于不同的模型, 为此DLK内置了很多有用的组件。

除了基础组件之外，DLK还为主要的任务提供了丰富的示例，更多的示例会慢慢添加进来

You Can Edit Your DLK Config Like Python Code

基于intc 所提供的强大的Config管理能力, 你可以像编写python代码一样编写你的config文件

Install

pip install dlk == 0.1.0

# or clone this repo and cd to the project root dir
pip install .

Demo

下面是一些基于dlk开发的示例:

NOTE: 由于我目前只有一台拥有一张AMD Radeon VII 16G的GPU和32G内存的个人PC，算力十分有限，因此这里示例的参数很多都还没有优化至SOTA

Grid Search

dlk基于intc进行开发，因此同样提供了参数搜索的能力，而intc的_search并不仅限于数值类型的参数搜索，也可以对整个模块进行搜索，因此dlk实际上也具有模块级的架构搜索能力

./examples/grid_search_exp里面提供了一个对超参数进行搜索的示例

训练完模型之后执行：

tensorboard --logdir ./logs

Task Demo

Demo 均位于examples目录下，训练完模型后执行：

streamlit run ./demo.py

Usage and Feature

使用方法

一般来说一个常见的dlk开发任务包含两个pipeline，即数据预处理pipeline和模型推理pipeline. 实际上这两个步骤是可以放到同一个pipeline中的, 当前示例中的大多数任务都需要对预处理数据的复用，因此使用两个pipeline

数据预处理pipeline对应的内置入口是dlk.preprocess.Process，我们需要编写process.jsonc config文件来对预处理过程(训练、推理、deploy过程全都复用同一个文件，因此配置文件中有针对不同stage的不同设置)进行配置并初始化Process, 将数据传入并执行run即可按要求输出预处理好的数据

模型训练pipeline对应的内置入口是dlk.train.Train, 我们需要编写fit.jsonc config文件来对模型训练（推理和deploy过程也同样复用这个文件），使用配置文件初始化Train之后执行run即可获得训练好的模型。

demo则只需要导入训练过程中相同的process.jsonc和fit.jsonc以及训练好的模型（由checkpoint callback组件保存）即可

模型部署只需将dlk.server.Server实例化，分发到对应的服务器，通过Server.fit接受单条或一个batch的数据即可（TODO: 示例）

模块注册

DLK依赖两个注册系统，一套是intc的config注册cregister，一套是dlk自己的模块注册，注册原则是一致的，都是将一个模块以module_type和module_name为key注册到注册器中，之所以选择两层的命名作为key是因为这样更方便区分不同的模块类型

以dlk.nn.layer.embedding.static 为例，我们将StaticEmbeddingConfig 作为StaticEmbedding的config以("embedding", "static")为key注册到intc的cregister中，以同样的key将StaticEmbedding注册到dlk的模块注册器register中。

使用注册器的好处是，我们可以不必关注具体类在哪里实现，只要知道注册的名称就可以直接获取这个类，这使得我们可以非常方便的在任意位置扩展embedding的类型，对于我们在自己的项目里面扩展dlk非常重要，注册模块对于intc也同样重要。在我们已知StaticEmbedding的注册名的情况下，获取这个模块的方法非常简单，可以直接register.get("embedding", "static")即可，而不必关注他的实际存储位置(cregister也有同样的功能)

部分内置模块介绍

callback

dlk的Trainer是基于lightning.Trainer实现的，因此dlk同样可以使用lightning提供的callback, dlk.callback中包含一些常用的callback

虚拟对抗训练

Adversarial Training是一种常见的提升模型效果的技巧，dlk内置了一些常用的针对embedding的adv方法(dlk.adv_method)，./examples/adv_exp是一个使用示例

复杂训练控制

dlk的dlk.scheduler模块提供了多种的训练scheduler， dlk.nn.loss模块中的multi_loss同样针对多个loss提供了自由控制各种loss的能力

文本生成

dlk还参考fairseq的实现，实现了多种的token_sample方法，为文本生成提供非常强大的控制能力

实现你自己的模型

参考./examples/001_first_example 实现你自己的模型

看完例子之后。但你可能会有疑问，这似乎并不比我直接实现一个模型简单，甚至有很多概念让我觉得这看起来更复杂。 是的，如果你只是想训练一个简单的模型，不需要考虑预测、演示等，没错，但是dlk提供了一个非常统一的框架，让你只需按照步骤来实现相应的组件，就可以获得一个可用的模型。并且所有的工作都是可重用的，包括你刚刚实现的组件。

而且dlk还提供了很多优化方面的工具，让你不是止步于简单模型

记住这个包的原则是Donot Repeat Yourself

More Document

TODO