Make Mindspore Training Easier

These details have not been verified by PyPI

Project links

Homepage

GitHub Statistics

View statistics for this project via Libraries.io, or by using our public dataset on Google BigQuery

Project description

MindsporeTrainer

PyPI PyPI - Python Version GitHub Workflow Status issues License

MindsporeTrainer 是基于昇思MindSpore的训练框架。让Mindspore的算法研究更容易一些。

Mindspore上手不易，希望能帮助各位炼丹师的升级之路更容易些。

Home page

主要特性

主要的几个出发点是：

采用纯python实现，方便多卡训练过程中的调试
易于扩展，对新任务采用插件式接入
方便实现多种模型的训练、评估、预测等

安装

pip

pip install MindsporeTrainer

源码

python setup.py

使用方法

DeBERTa预训练任务示例

1. 定义task

MindsporeTrainer/apps/tasks/deberta.py

from collections import OrderedDict
import numpy as np
import os
import random
from shutil import copyfile
from loguru import logger

from mindspore.communication import get_rank, get_group_size

from MindsporeTrainer.data import ExampleInstance, ExampleSet
from MindsporeTrainer.data.example import *
from MindsporeTrainer.apps.tasks import EvalData, TransformerTask
from MindsporeTrainer.apps.tasks import register_task
from MindsporeTrainer.utils.metrics import *
from MindsporeTrainer.utils.metrics import BertMetric
from MindsporeTrainer.utils.masker import NGramMaskGenerator
from MindsporeTrainer.data.dynamic_dataset import create_dynamic_dataset
from MindsporeTrainer.modeling.tokenizers import BertTokenizer

@register_task(name="DEBERTA", desc="Basic DEBERTA task")
class DEBERTATask(TransformerTask):
    def __init__(self, data_dir, args, **kwargs):
        super().__init__(args, **kwargs)
        self.max_seq_length = 512
        self.model_config = 'data/pretrained_models/deberta-base-v2/model_config.json'
        self.vocab_type = 'BERT'
        self.vocab_path = 'data/pretrained_models/deberta-base-v2/vocab.txt'
        self.data_dir = data_dir
        self.args = args
        self.metric = 'bert'
        self.main_metric = 'perplexity'
        self.optimizer_name = 'Lamb'
        self.tokenizer = BertTokenizer(self.vocab_path)
        if args.distribute:
            self.rank_id = get_rank()
            self.rank_size = get_group_size()
        else:
            self.rank_id = 0
            self.rank_size = 1

    def train_data(self, max_seq_len=512, batch_size=32, **kwargs):
        data_path = os.path.join(self.data_dir, 'daizhige.pkl')
        data = self.load_data(data_path, 'GW', max_seq_len)
        # data = ExampleSet(data)
        output_columns = ["input_ids", "input_mask", "token_type_id", "next_sentence_labels",
                                    "masked_lm_positions", "masked_lm_ids", "masked_lm_weights"]
        return create_dynamic_dataset(data, self.get_feature_fn(), 
                                      batch_size=batch_size,
                                      output_columns=output_columns, 
                                      repeat=self.args.num_train_epochs,
                                      num_workers=self.args.data_workers,
                                      num_shards=self.rank_size, 
                                      shard_id=self.rank_id)

    def eval_data(self, max_seq_len=512, batch_size=32, **kwargs):
        ...

    def get_metrics(self, **kwargs):
        """Calcuate metrics based on prediction results"""
        return OrderedDict(
            bert = BertMetric(self.args.eval_batch_size),
            )

    def get_eval_fn(self, **kwargs):
        data = kwargs.get('data')
        if data is None:
            data = self.eval_data(**kwargs)
        def run_eval(model, name, prefix):
            '''
            args: 
                model: Model instance
                name: evaluate name
                prefix: prefix of file
            return:
                float, main metric of this task, used to save best metric model
            '''
            res = model.eval(data, dataset_sink_mode=False)
            res = res['bert']
            main_metric = res[self.main_metric]
            if self.rank_id == 0:
                output=os.path.join(self.args.output_dir, 'submit-{}-{}.tsv'.format(name, prefix))
                metric_str = '\n'.join([f'{k}: {v:.4f}' for k, v in res.items()])
                metric_str = metric_str + '\n'
                logger.info("====================================")
                logger.info("evaluate result:\n")
                logger.info(metric_str)
                logger.info("====================================")

                with open(output, 'w', encoding='utf-8') as fs:
                    fs.write(f"metrics:\n{metric_str}\n")
            return main_metric
        return run_eval


    def get_feature_fn(self, max_seq_len=512, ext_params=None, rng=None, **kwargs):
        tokenizer = self.tokenizer
        mask_generator = NGramMaskGenerator(tokenizer)
        def example_to_feature(*example):
            '''
            sample: text, label
            return: ["input_ids", "input_mask", "token_type_id", "next_sentence_labels",
                    masked_lm_positions", "masked_lm_ids", "masked_lm_weights"]
            '''
            ......
            return tuple(features)

        return example_to_feature

    def load_data(self, path, type=None, max_seq_len=512):
        examples = []
        ......
        return ExampleSet(examples)


    def get_model(self):
        from MindsporeTrainer.modeling.models import BertPreTraining, DebertaPreTraining
        from MindsporeTrainer import build_transformer_model
        if self.args.fp16:
            compute_type = mstype.float16
        else:
            compute_type = mstype.float32
        model, config = build_transformer_model(self.model_config, 
                                                model=DebertaPreTraining, 
                                                compute_type=compute_type, 
                                                padding_idx=self.tokenizer._convert_token_to_id(self.tokenizer.pad_token))
        with open(os.path.join(self.args.output_dir, 'config.json'), 'w', encoding='utf-8') as f:
            f.write(config.to_json_string())
        copyfile(self.vocab_path, os.path.join(self.args.output_dir, 'vocab.txt'))
        return model
        # return partial_class

    def get_loss(self, *args, **kwargs):
        from MindsporeTrainer.modeling.loss import BertPretrainingLoss
        return BertPretrainingLoss(self.tokenizer.vocab_size)

    def get_eval_head(self, *args, **kwargs):
        from MindsporeTrainer.modeling.models import BertEvalHead
        return BertEvalHead(self.tokenizer.vocab_size)

    def get_opt_fn(self, *args, **kwargs):
        return None

2. 编写启动脚本

run.py

import MindsporeTrainer as mst
mst.launch()

3. 运行任务

python run.py --task_name=RESNETTask --do_train --do_eval --data_dir=data --num_train_epochs=10 --learning_rate=1e-2 --train_batch_size=64 --save_eval_steps=1000 --output_dir=output

多卡训练

官方推荐编写bash脚本，利用mpi进行训练，这里采用了纯python的实现。

定义必须的环境变量

bash

export RANK_SIZE = 8
export RANK_TABLE_FILE = /path/hccl.json

vscode调试环境

"env": {
    "RANK_SIZE": "8",
    "RANK_TABLE_FILE": "/path/hccl.json"
}

设置参数，开始训练

python run.py ...... --device_num=8 --device_id=0,1,2,3,4,5,6,7

模型创建

自定义模型

Mindspore相对TF和PyTorch，有其自身的特点，建模习惯也不同。在本框架中，为了便于模块化与代码复用，将模型分为backbone和predict head两部分。
需要注意的是，模型定义中，construct函数返回都应当是tuple，即使只有一个对象返回，也应当采用(obj,)的形式。

backbone

模型的主体部分定义。

predict head

模型的附属部分，根据模型的用途，通常可以定义为loss层，或者在evaluation过程中定义的eval head，又或者其他用途的头。

使用官方model_zoo中的模型

其代码结构如下

models
├── official                                    # 官方支持模型
│   └── XXX                                     # 模型名
│       ├── README.md                           # 模型说明文档
│       ├── requirements.txt                    # 依赖说明文件
│       ├── eval.py                             # 精度验证脚本
│       ├── export.py                           # 推理模型导出脚本
│       ├── scripts                             # 脚本文件
│       │   ├── run_distributed_train.sh        # 分布式训练脚本
│       │   ├── run_eval.sh                     # 验证脚本
│       │   └── run_standalone_train.sh         # 单机训练脚本
│       ├── src                                 # 模型定义源码目录
│       │   ├── XXXNet.py                       # 模型结构定义
│       │   ├── callback.py                     # 回调函数定义
│       │   ├── config.py                       # 模型配置参数文件
│       │   └── dataset.py                      # 数据集处理定义
│       ├── ascend_infer                        # （可选）用于在Ascend推理设备上进行离线推理的脚本
│       ├── third_party                         # （可选）第三方代码
│       │   └── XXXrepo                         # （可选）完整克隆自第三方仓库的代码
│       └── train.py                            # 训练脚本
├── research                                    # 非官方研究脚本
├── community                                   # 合作方脚本链接
└── utils                                       # 模型通用工具

找到所需的模型目录后，在src/xxxmodel.py中找到相应的定义，如果有定义好的backbone和head，那么可以直接引入使用。
例如使用其中的BERT模型：
git clone https://gitee.com/mindspore/models
复制出bert源码到工作目录models/official/nlp/bert/src
定义task

    from bert.src.bert_for_pre_training import BertPreTraining, BertPretrainingLoss
    ......
    def get_model(self):
        from MindsporeTrainer import build_transformer_model
        if self.args.fp16:
            compute_type = mstype.float16
        else:
            compute_type = mstype.float32
        model, config = build_transformer_model(self.model_config, 
                                                model=BertPreTraining, 
                                                compute_type=compute_type, 
                                                padding_idx=self.tokenizer._convert_token_to_id(self.tokenizer.pad_token))
        with open(os.path.join(self.output_dir, 'config.json'), 'w', encoding='utf-8') as f:
            f.write(config.to_json_string())
        copyfile(self.vocab_path, os.path.join(self.output_dir, 'vocab.txt'))
        return model

    def get_loss(self, *args, **kwargs):
        return BertPretrainingLoss(self.tokenizer.vocab_size)

    def get_eval_head(self, *args, **kwargs):
        from MindsporeTrainer.modeling.models import BertEvalHead
        return BertEvalHead(self.tokenizer.vocab_size)

参数介绍

训练超参数基本上都是通过运行参数来控制的，另外一些则可以在task定义中控制。

commmon args

--accumulation_steps
type=int
default=1
Accumulating gradients N times before weight update, default is 1.
--allreduce_post_accumulation
type=str
default=true
choices=[true, false]
Whether to allreduce after accumulation of N steps or after each step, default is true.
--data_dir
default=None
type=str
required=False
The input data dir. Should contain the .tsv files (or other data files) for the task.
--data_sink_steps
type=int
default=1
Sink steps for each epoch, default is 1.
--do_train
default=False
action='store_true'
Whether to run training.
--do_eval
default=False
action='store_true'
Whether to run eval on the dev set.
--do_predict
default=False
action='store_true'
Whether to run prediction on the test set.
--debug
default=False
action='store_true'
Whether to cache cooked binary features
--device_target
type=str
default='Ascend'
choices=['Ascend 'GPU']
device where the code will be implemented. (Default: Ascend)
--distribute
default=False
action='store_true'
Run distribute, default is false.
--device_id
type=str
default=0
Device id, default is 0.
--device_num
type=int
default=1
Use device nums, default is 1.
--enable_data_sink
default=False
action='store_true'
Enable data sink, default is false.
--load_checkpoint_path
type=str
default=''
Load checkpoint file path
--num_train_epochs
default=1
type=int
Total number of training epochs to perform.
--output_dir
default=None
type=str
required=True
The output directory where the model checkpoints will be written.
--run_mode
type=str
default='GRAPH'
choices=['GRAPH 'PY']
0: GRAPH_MODE, 1: PY_NATIVE_MODE
--save_eval_steps
type=int
default=1000
Save checkpoint and evaluate steps, default is 1000.
--save_checkpoint_num
type=int
default=1
Save checkpoint numbers, default is 1.
--tag
type=str
default='final'
The tag name of current prediction/runs.
--task_dir
default=None
type=str
required=False
The directory to load customized tasks.
--task_name
default=None
type=str
action=LoadTaskAction
required=True
The name of the task to train.

train args

--data_workers
type=int
default=4
The workers to load data.
--enable_graph_kernel
type=str
default=auto
choices=[auto, true, false]
Accelerate by graph kernel, default is auto.
--eval_batch_size
default=32
type=int
Total batch size for eval.
--enable_global_norm
type=bool
default=False
enable global norm
--predict_batch_size
default=32
type=int
Total batch size for prediction.
--report_interval
default=1
type=int
Interval steps for state report.
--save_graphs
default=False
action='store_true'
Whether to save graphs
--seed
type=int
default=1234
random seed for initialization
--thor
default=False
action='store_true' Whether to convert model to thor optimizer
--train_batch_size
default=64
type=int
Total batch size for training.
--train_steps
type=int
default=-1
Training Steps, default is -1, meaning run all steps according to epoch number.

optimizer args

--fp16
default=False
type=boolean_string
Whether to use 16-bit float precision instead of 32-bit
--learning_rate
default=5e-5
type=float
The initial learning rate for Adam.
--loss_scale_value
type=int
default=1024
initial loss scale value
--resume_opt_path
type=str.lower
default=''
The optimizer to be resumed.
--scale_factor
type=int
default=4
loss scale factor
--scale_window
type=int
default=1000
loss window
--warmup
default=0.1
type=float
Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10% of training.

task中增加自定义参数

    @classmethod
    def add_arguments(cls, parser):
        """Add task specific arguments
            e.g. parser.add_argument('--data_dir', type=str, help='The path of data directory.')
        """
        parser.add_argument('--task_example_arg', type=str, default=None, help='An example task specific argument')

        return parser

Task

所有的task都应当继承于MindsporeTrainer.apps.tasks.Task,
为transformer定义的MindsporeTrainer.apps.tasks.TransformerTask也继承自Task。
Task类的定义为:

class Task():
    _meta={}

    def __init__(self, args, **kwargs):
        self.args = args
    
    def eval_data(self, **kwargs):
        """
        Get eval dataset object.
        """
        return None

    def train_data(self, **kwargs):
        """
        Get train dataset object.
        """
        return None

    def test_data(self, **kwargs):
        return None

    def get_labels(self):
        """Gets the list of labels for this data set."""
        return None

    def get_eval_fn(self, *args, **kwargs):
        """
        Get the evaluation function
        """
        return None

    def get_eval_head(self, *args, **kwargs):
        """
        Get the evaluate head, the head replace loss function head when evaluation process
        """
        return None

    def get_pred_fn(self, *args, **kwargs):
        """
        Get the predict function
        """
        return None

    def get_loss(self, *args, **kwargs):
        """
        Get the loss function
        """
        return None

    def get_opt_fn(self, *args, **kwargs):
        """
        Get a function wich return the opimizer
        """
        def get_optimizer(*args, **kwargs):
            pass
        return get_optimizer

    def get_metrics(self):
        """Calcuate metrics based on prediction results"""
        return None

    def get_predict_fn(self):
        """Calcuate metrics based on prediction results"""
        def predict_fn(logits, output_dir, name, prefix):
            pass
        return None

    def get_feature_fn(self, **kwargs):
        """
        get the featurize function
        """
        def _example_to_feature(**kwargs):
             return feature
        return _example_to_feature
    
    def get_model(self):
        """
        Get a model instance
        """
        raise NotImplementedError('method not implemented yet.')

    @classmethod
    def add_arguments(cls, parser):
        """Add task specific arguments
        """
        pass

API