Skip to content
Snippets Groups Projects
Unverified Commit 01f8c6a7 authored by i-robot's avatar i-robot Committed by Gitee
Browse files

!2320 remove remained quant models

Merge pull request !2320 from cjh9368/master
parents c7095b79 7b32ec2c
No related branches found
No related tags found
No related merge requests found
Showing
with 0 additions and 3133 deletions
ARG FROM_IMAGE_NAME
FROM ${FROM_IMAGE_NAME}
COPY requirements.txt .
RUN pip3.7 install -r requirements.txt
\ No newline at end of file
#!/bin/bash
# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.mitations under the License.
docker_image=$1
data_dir=$2
model_dir=$3
docker run -it --ipc=host \
--device=/dev/davinci0 \
--device=/dev/davinci1 \
--device=/dev/davinci2 \
--device=/dev/davinci3 \
--device=/dev/davinci4 \
--device=/dev/davinci5 \
--device=/dev/davinci6 \
--device=/dev/davinci7 \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm --device=/dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/Ascend/add-ons/:/usr/local/Ascend/add-ons/ \
-v ${model_dir}:${model_dir} \
-v ${data_dir}:${data_dir} \
-v /root/ascend/log:/root/ascend/log ${docker_image} /bin/bash
\ No newline at end of file
# Contents
- [Contents](#contents)
- [TernaryBERT Description](#ternarybert-description)
- [Model Architecture](#model-architecture)
- [Dataset](#dataset)
- [Environment Requirements](#environment-requirements)
- [Quick Start](#quick-start)
- [Script Description](#script-description)
- [Script and Sample Code](#script-and-sample-code)
- [Script Parameters](#script-parameters)
- [Train](#train)
- [Eval](#eval)
- [Options and Parameters](#options-and-parameters)
- [Parameters](#parameters)
- [Training Process](#training-process)
- [Training](#training)
- [Evaluation Process](#evaluation-process)
- [Evaluation](#evaluation)
- [evaluation on STS-B dataset](#evaluation-on-STS-B-dataset)
- [Model Description](#model-description)
- [Performance](#performance)
- [training Performance](#training-performance)
- [Inference Performance](#inference-performance)
- [Description of Random Situation](#description-of-random-situation)
- [ModelZoo Homepage](#modelzoo-homepage)
# [TernaryBERT Description](#contents)
[TernaryBERT](https://arxiv.org/abs/2009.12812) ternarizes the weights in a fine-tuned [BERT](https://arxiv.org/abs/1810.04805) or [TinyBERT](https://arxiv.org/abs/1909.10351) model and achieves competitive performances in natural language processing tasks. TernaryBERT outperforms the other BERT quantization methods, and even achieves comparable performance as the full-precision model while being 14.9x smaller
[Paper](https://arxiv.org/abs/2009.12812): Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang and Qun Liu. [TernaryBERT: Distillation-aware Ultra-low Bit BERT](https://arxiv.org/abs/2009.12812). arXiv preprint arXiv:2009.12812.
# [Model Architecture](#contents)
The backbone structure of TernaryBERT is transformer, the transformer contains six encoder modules, one encoder contains one self-attention module and one self-attention module contains one attention module. The pretrained teacher model and student model are provided [here](https://download.mindspore.cn/model_zoo/research/nlp/ternarybert/).
# [Dataset](#contents)
- Download glue dataset for task distillation. Convert dataset files from json format to tfrecord format, please refer to run_classifier.py which in [BERT](https://github.com/google-research/bert) repository.
# [Environment Requirements](#contents)
- Hardware(GPU)
- Prepare hardware environment with GPU processor.
- Framework
- [MindSpore](https://gitee.com/mindspore/mindspore)
- For more information, please check the resources below:
- [MindSpore Tutorials](https://www.mindspore.cn/tutorials/en/master/index.html)
- [MindSpore Python API](https://www.mindspore.cn/docs/api/en/master/index.html)
- Software:
- sklearn
# [Quick Start](#contents)
After installing MindSpore via the official website, you can start training and evaluation as follows:
```bash
# run training example
bash scripts/run_train.sh [TASK_NAME] [DEVICE_TARGET] [TEACHER_MODEL_DIR] [STUDENT_MODEL_DIR] [DATA_DIR] [DEVICE_ID](optional)
Before running the shell script, please set the `task_name`, `device_target`, `teacher_model_dir`, `student_model_dir`, `data_dir` and `device_id`(optional) in the run_train.sh file first.
# run evaluation example
bash scripts/run_eval.sh [TASK_NAME] [DEVICE_TARGET] [MODEL_DIR] [DATA_DIR]
Before running the shell script, please set the `task_name`, `device_target`, `model_dir` and `data_dir` in the run_eval.sh file first.
```
# [Script Description](#contents)
## [Script and Sample Code](#contents)
```text
.
└─ternarybert
├─README.md
├─scripts
├─run_train.sh # shell script for training phase
├─run_eval.sh # shell script for evaluation phase
├─src
├─__init__.py
├─assessment_method.py # assessment method for evaluation
├─cell_wrapper.py # cell for training
├─config.py # parameter configuration for training and evaluation phase
├─dataset.py # data processing
├─quant.py # function for quantization
├─tinybert_model.py # backbone code of network
├─utils.py # util function
├─__init__.py
├─train.py # train net for task distillation
├─eval.py # evaluate net after task distillation
```
## [Script Parameters](#contents)
### Train
```text
usage: train.py [--h] [--device_target GPU] [--do_eval {true,false}] [--epoch_size EPOCH_SIZE]
[--device_id DEVICE_ID] [--do_shuffle {true,false}] [--enable_data_sink {true,false}] [--save_ckpt_step SAVE_CKPT_STEP]
[--eval_ckpt_step EVAL_CKPT_STEP] [--max_ckpt_num MAX_CKPT_NUM] [--data_sink_steps DATA_SINK_STEPS]
[--teacher_model_dir TEACHER_MODEL_DIR] [--student_model_dir STUDENT_MODEL_DIR] [--data_dir DATA_DIR]
[--output_dir OUTPUT_DIR] [--task_name {sts-b,qnli,mnli}] [--dataset_type DATASET_TYPE] [--seed SEED]
[--train_batch_size TRAIN_BATCH_SIZE] [--eval_batch_size EVAL_BATCH_SIZE]
options:
--device_target Device where the code will be implemented: "GPU"
--do_eval Do eval task during training or not: "true" | "false", default is "true"
--epoch_size Epoch size for train phase: N, default is 3
--device_id Device id: N, default is 0
--do_shuffle Enable shuffle for train dataset: "true" | "false", default is "true"
--enable_data_sink Enable data sink: "true" | "false", default is "true"
--save_ckpt_step If do_eval is false, the checkpoint will be saved every save_ckpt_step: N, default is 50
--eval_ckpt_step If do_eval is true, the evaluation will be ran every eval_ckpt_step: N, default is 50
--max_ckpt_num The number of checkpoints will not be larger than max_ckpt_num: N, default is 50
--data_sink_steps Sink steps for each epoch: N, default is 1
--teacher_model_dir The checkpoint directory of teacher model: PATH, default is ""
--student_model_dir The checkpoint directory of student model: PATH, default is ""
--data_dir Data directory: PATH, default is ""
--output_dir The output checkpoint directory: PATH, default is "./"
--task_name The name of the task to train: "sts-b" | "qnli" | "mnli", default is "sts-b"
--dataset_type The name of the task to train: "tfrecord" | "mindrecord", default is "tfrecord"
--seed The random seed: N, default is 1
--train_batch_size Batch size for training: N, default is 16
--eval_batch_size Eval Batch size in callback: N, default is 32
```
### Eval
```text
usage: eval.py [--h] [--device_target GPU] [--device_id DEVICE_ID] [--model_dir MODEL_DIR] [--data_dir DATA_DIR]
[--task_name {sts-b,qnli,mnli}] [--dataset_type DATASET_TYPE] [--batch_size BATCH_SIZE]
options:
--device_target Device where the code will be implemented: "GPU"
--device_id Device id: N, default is 0
--model_dir The checkpoint directory of model: PATH, default is ""
--data_dir Data directory: PATH, default is ""
--task_name The name of the task to train: "sts-b" | "qnli" | "mnli", default is "sts-b"
--dataset_type The name of the task to train: "tfrecord" | "mindrecord", default is "tfrecord"
--batch_size Batch size for evaluating: N, default is 32
```
## Parameters
`config.py`contains parameters of glue tasks, train, optimizer, eval, teacher BERT model and student BERT model.
```text
Parameters for glue task:
num_labels the numbers of labels: N.
seq_length length of input sequence: N
task_type the type of task: "classification" | "regression"
metrics the eval metric for task: Accuracy | F1 | Pearsonr | Matthews
Parameters for train:
batch_size batch size of input dataset: N, default is 16
loss_scale_value initial value of loss scale: N, default is 2^16
scale_factor factor used to update loss scale: N, default is 2
scale_window steps for once updatation of loss scale: N, default is 50
Parameters for optimizer:
learning_rate value of learning rate: Q, default is 5e-5
end_learning_rate value of end learning rate: Q, must be positive, default is 1e-14
power power: Q, default is 1.0
weight_decay weight decay: Q, default is 1e-4
eps term added to the denominator to improve numerical stability: Q, default is 1e-6
warmup_ratio the ratio of warmup steps to total steps: Q, default is 0.1
Parameters for eval:
batch_size batch size of input dataset: N, default is 32
Parameters for teacher bert network:
seq_length length of input sequence: N, default is 128
vocab_size size of each embedding vector: N, must be consistent with the dataset you use. Default is 30522
hidden_size size of bert encoder layers: N
num_hidden_layers number of hidden layers: N
num_attention_heads number of attention heads: N, default is 12
intermediate_size size of intermediate layer: N
hidden_act activation function used: ACTIVATION, default is "gelu"
hidden_dropout_prob dropout probability for BertOutput: Q
attention_probs_dropout_prob dropout probability for BertAttention: Q
max_position_embeddings maximum length of sequences: N, default is 512
save_ckpt_step number for saving checkponit: N, default is 100
max_ckpt_num maximum number for saving checkpoint: N, default is 1
type_vocab_size size of token type vocab: N, default is 2
initializer_range initialization value of Normal: Q, default is 0.02
use_relative_positions use relative positions or not: True | False, default is False
dtype data type of input: mstype.float16 | mstype.float32, default is mstype.float32
compute_type compute type in BertTransformer: mstype.float16 | mstype.float32, default is mstype.float32
Parameters for student bert network:
seq_length length of input sequence: N, default is 128
vocab_size size of each embedding vector: N, must be consistent with the dataset you use. Default is 30522
hidden_size size of bert encoder layers: N
num_hidden_layers number of hidden layers: N
num_attention_heads number of attention heads: N, default is 12
intermediate_size size of intermediate layer: N
hidden_act activation function used: ACTIVATION, default is "gelu"
hidden_dropout_prob dropout probability for BertOutput: Q
attention_probs_dropout_prob dropout probability for BertAttention: Q
max_position_embeddings maximum length of sequences: N, default is 512
save_ckpt_step number for saving checkponit: N, default is 100
max_ckpt_num maximum number for saving checkpoint: N, default is 1
type_vocab_size size of token type vocab: N, default is 2
initializer_range initialization value of Normal: Q, default is 0.02
use_relative_positions use relative positions or not: True | False, default is False
dtype data type of input: mstype.float16 | mstype.float32, default is mstype.float32
compute_type compute type in BertTransformer: mstype.float16 | mstype.float32, default is mstype.float32
do_quant do activation quantilization or not: True | False, default is True
embedding_bits the quant bits of embedding: N, default is 2
weight_bits the quant bits of weight: N, default is 2
cls_dropout_prob dropout probability for BertModelCLS: Q
activation_init initialization value of activation quantilization: Q, default is 2.5
is_lgt_fit use label ground truth loss or not: True | False, default is False
```
## [Training Process](#contents)
### Training
Before running the command below, please check `teacher_model_dir`, `student_model_dir` and `data_dir` has been set. Please set the path to be the absolute full path, e.g:"/home/xxx/model_dir/".
```text
python
python train.py --task_name='sts-b' --device_target="GPU" --teacher_model_dir='/home/xxx/model_dir/' --student_model_dir='/home/xxx/model_dir/' --data_dir='/home/xxx/data_dir/'
shell
bash scripts/run_train.sh [TASK_NAME] [DEVICE_TARGET] [TEACHER_MODEL_DIR] [STUDENT_MODEL_DIR] [DATA_DIR] [DEVICE_ID](optional)
```
The shell command above will run in the background, you can view the results the file log.txt. The python command will run in the console, you can view the results on the interface. After training, you will get some checkpoint files under the script folder by default. The eval metric value will be achieved as follows:
```text
step: 50, Pearsonr 72.50008506516072, best_Pearsonr 72.50008506516072
step 100, Pearsonr 81.3580301181608, best_Pearsonr 81.3580301181608
step 150, Pearsonr 83.60461724688754, best_Pearsonr 83.60461724688754
step 200, Pearsonr 82.23210161651377, best_Pearsonr 83.60461724688754
...
step 1050, Pearsonr 87.5606067964618332, best_Pearsonr 87.58388835685436
```
## [Evaluation Process](#contents)
### Evaluation
If you want to after running and continue to eval.
#### evaluation on STS-B dataset
```text
python
python eval.py --task_name='sts-b' --device_target="GPU" --model_dir='/home/xxx/model_dir/' --data_dir='/home/xxx/data_dir/'
shell
bash scripts/run_eval.sh [TASK_NAME] [DEVICE_TARGET] [MODEL_DIR] [DATA_DIR]
```
The shell command above will run in the background, you can view the results the file log.txt. The python command will run in the console, you can view the results on the interface. The metric value of the test dataset will be as follows:
```text
eval step: 0, Pearsonr: 96.91109003302263
eval step: 1, Pearsonr: 95.6800637493701
eval step: 2, Pearsonr: 94.23823082886167
...
The best Pearsonr: 87.58388835685437
```
## [Model Description](#contents)
## [Performance](#contents)
### training Performance
| Parameters | GPU |
| -------------------------- | ------------------------- |
| Model Version | TernaryBERT |
| Resource | NV GeForce GTX1080ti, cpu 2.00GHz, cores 56, mem 251G, os Ubuntu 16.04 |
| Date | 2021-6-10 |
| MindSpore Version | 1.2.0 |
| Dataset | STS-B |
| batch_size | 16 |
| Metric value | 87.5 |
# [Description of Random Situation](#contents)
In train.py, we set do_shuffle to shuffle the dataset.
In config.py, we set the hidden_dropout_prob, attention_pros_dropout_prob and cls_dropout_prob to dropout some network node.
# [ModelZoo Homepage](#contents)
Please check the official [homepage](https://gitee.com/mindspore/models).
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""eval standalone script"""
import os
import re
import argparse
from mindspore import context
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from src.dataset import create_dataset
from src.config import eval_cfg, student_net_cfg, task_cfg
from src.tinybert_model import BertModelCLS
DATA_NAME = 'eval.tf_record'
def parse_args():
"""
parse args
"""
parser = argparse.ArgumentParser(description='ternarybert evaluation')
parser.add_argument('--device_target', type=str, default='GPU', choices=['Ascend', 'GPU'],
help='Device where the code will be implemented. (Default: GPU)')
parser.add_argument('--device_id', type=int, default=0, help='Device id. (Default: 0)')
parser.add_argument('--model_dir', type=str, default='', help='The checkpoint directory of model.')
parser.add_argument('--data_dir', type=str, default='', help='Data directory.')
parser.add_argument('--task_name', type=str, default='sts-b', choices=['sts-b', 'qnli', 'mnli'],
help='The name of the task to train. (Default: sts-b)')
parser.add_argument('--dataset_type', type=str, default='tfrecord', choices=['tfrecord', 'mindrecord'],
help='The name of the task to train. (Default: tfrecord)')
parser.add_argument('--batch_size', type=int, default=32, help='Batch size for evaluating')
return parser.parse_args()
def get_ckpt(ckpt_file):
lists = os.listdir(ckpt_file)
lists.sort(key=lambda fn: os.path.getmtime(ckpt_file + '/' + fn))
return os.path.join(ckpt_file, lists[-1])
def do_eval_standalone(args_opt):
"""
do eval standalone
"""
ckpt_file = os.path.join(args_opt.model_dir, args_opt.task_name)
ckpt_file = get_ckpt(ckpt_file)
print('ckpt file:', ckpt_file)
task = task_cfg[args_opt.task_name]
student_net_cfg.seq_length = task.seq_length
eval_cfg.batch_size = args_opt.batch_size
eval_data_dir = os.path.join(args_opt.data_dir, args_opt.task_name, DATA_NAME)
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args.device_id)
eval_dataset = create_dataset(batch_size=eval_cfg.batch_size,
device_num=1,
rank=0,
do_shuffle=False,
data_dir=eval_data_dir,
data_type=args_opt.dataset_type,
seq_length=task.seq_length,
task_type=task.task_type,
drop_remainder=False)
print('eval dataset size:', eval_dataset.get_dataset_size())
print('eval dataset batch size:', eval_dataset.get_batch_size())
eval_model = BertModelCLS(student_net_cfg, False, task.num_labels, 0.0, phase_type='student')
param_dict = load_checkpoint(ckpt_file)
new_param_dict = {}
for key, value in param_dict.items():
new_key = re.sub('tinybert_', 'bert_', key)
new_key = re.sub('^bert.', '', new_key)
new_param_dict[new_key] = value
load_param_into_net(eval_model, new_param_dict)
eval_model.set_train(False)
columns_list = ["input_ids", "input_mask", "segment_ids", "label_ids"]
callback = task.metrics()
for step, data in enumerate(eval_dataset.create_dict_iterator()):
input_data = []
for i in columns_list:
input_data.append(data[i])
input_ids, input_mask, token_type_id, label_ids = input_data
_, _, logits, _ = eval_model(input_ids, token_type_id, input_mask)
callback.update(logits, label_ids)
print('eval step: {}, {}: {}'.format(step, callback.name, callback.get_metrics()))
metrics = callback.get_metrics()
print('The best {}: {}'.format(callback.name, metrics))
if __name__ == '__main__':
args = parse_args()
do_eval_standalone(args)
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Bert hub interface for bert base"""
from src.tinybert_model import BertModel
from src.tinybert_model import BertConfig
import mindspore.common.dtype as mstype
tinybert_student_net_cfg = BertConfig(
seq_length=128,
vocab_size=30522,
hidden_size=768,
num_hidden_layers=6,
num_attention_heads=12,
intermediate_size=3072,
hidden_act="gelu",
hidden_dropout_prob=0.1,
attention_probs_dropout_prob=0.1,
max_position_embeddings=512,
type_vocab_size=2,
initializer_range=0.02,
use_relative_positions=False,
dtype=mstype.float32,
compute_type=mstype.float32,
do_quant=True,
embedding_bits=2,
weight_bits=2,
weight_clip_value=3.0,
cls_dropout_prob=0.1,
activation_init=2.5,
is_lgt_fit=False
)
def create_network(name, *args, **kwargs):
"""
Create tinybert network.
"""
if name == "ternarybert":
if "seq_length" in kwargs:
tinybert_student_net_cfg.seq_length = kwargs["seq_length"]
is_training = kwargs.get("is_training", False)
return BertModel(tinybert_student_net_cfg, is_training, *args)
raise NotImplementedError(f"{name} is not implemented in the repo")
numpy
easydict
#!/bin/bash
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [ $# != 4 ]
then
echo "============================================================================================================"
echo "Please run the script as: "
echo "sh scripts/run_eval.sh [TASK_NAME] [DEVICE_TARGET] [MODEL_DIR] [DATA_DIR]"
echo "============================================================================================================"
exit 1
fi
echo "===============================================start evaling================================================"
task_name=$1
device_target=$2
model_dir=$3
data_dir=$4
mkdir -p ms_log
PROJECT_DIR=$(cd "$(dirname "$0")" || exit; pwd)
CUR_DIR=`pwd`
export GLOG_log_dir=${CUR_DIR}/ms_log
export GLOG_logtostderr=0
python ${PROJECT_DIR}/../eval.py \
--task_name=$task_name \
--device_target=$device_target \
--device_id=0 \
--model_dir=$model_dir \
--data_dir=$data_dir > eval_log.txt 2>&1 &
#!/bin/bash
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
if [[ $# != 5 && $# != 6 ]]
then
echo "============================================================================================================"
echo "Please run the script as: "
echo "sh scripts/run_train.sh [TASK_NAME] [DEVICE_TARGET] [TEACHER_MODEL_DIR] [STUDENT_MODEL_DIR] [DATA_DIR] [DEVICE_ID]"
echo "DEVICE_ID is optional, it can be set by command line argument DEVICE_ID , otherwise the value is zero"
echo "============================================================================================================"
exit 1
fi
echo "===============================================start training==============================================="
task_name=$1
device_target=$2
teacher_model_dir=$3
student_model_dir=$4
data_dir=$5
if [ $# == 6 ]
then
device_id=$6
else
device_id=0
fi
if [ $device_target == "GPU" ]
then
export CUDA_VISIBLE_DEVICES=$device_id
fi
mkdir -p ms_log
PROJECT_DIR=$(cd "$(dirname "$0")" || exit; pwd)
CUR_DIR=`pwd`
export GLOG_log_dir=${CUR_DIR}/ms_log
export GLOG_logtostderr=0
python ${PROJECT_DIR}/../train.py \
--task_name=$task_name \
--device_target=$device_target \
--device_id=$device_id \
--teacher_model_dir=$teacher_model_dir \
--student_model_dir=$student_model_dir \
--data_dir=$data_dir> train_log.txt 2>&1 &
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""assessment methods"""
import numpy as np
class Accuracy:
"""Accuracy"""
def __init__(self):
self.acc_num = 0
self.total_num = 0
self.name = 'Accuracy'
def update(self, logits, labels):
labels = labels.asnumpy()
labels = np.reshape(labels, -1)
logits = logits.asnumpy()
logit_id = np.argmax(logits, axis=-1)
self.acc_num += np.sum(labels == logit_id)
self.total_num += len(labels)
def get_metrics(self):
return self.acc_num / self.total_num * 100.0
class F1:
"""F1"""
def __init__(self):
self.logits_array = np.array([])
self.labels_array = np.array([])
self.name = 'F1'
def update(self, logits, labels):
labels = labels.asnumpy()
labels = np.reshape(labels, -1)
logits = logits.asnumpy()
logits = np.argmax(logits, axis=1)
self.labels_array = np.concatenate([self.labels_array, labels]).astype(np.bool)
self.logits_array = np.concatenate([self.logits_array, logits]).astype(np.bool)
def get_metrics(self):
if len(self.labels_array) < 2:
return 0.0
tp = np.sum(self.labels_array & self.logits_array)
fp = np.sum(self.labels_array & (~self.logits_array))
fn = np.sum((~self.labels_array) & self.logits_array)
p = tp / (tp + fp)
r = tp / (tp + fn)
return 2.0 * p * r / (p + r) * 100.0
class Pearsonr:
"""Pearsonr"""
def __init__(self):
self.logits_array = np.array([])
self.labels_array = np.array([])
self.name = 'Pearsonr'
def update(self, logits, labels):
labels = labels.asnumpy()
labels = np.reshape(labels, -1)
logits = logits.asnumpy()
logits = np.reshape(logits, -1)
self.labels_array = np.concatenate([self.labels_array, labels])
self.logits_array = np.concatenate([self.logits_array, logits])
def get_metrics(self):
if len(self.labels_array) < 2:
return 0.0
x_mean = self.logits_array.mean()
y_mean = self.labels_array.mean()
xm = self.logits_array - x_mean
ym = self.labels_array - y_mean
norm_xm = np.linalg.norm(xm)
norm_ym = np.linalg.norm(ym)
return np.dot(xm / norm_xm, ym / norm_ym) * 100.0
class Matthews:
"""Matthews"""
def __init__(self):
self.logits_array = np.array([])
self.labels_array = np.array([])
self.name = 'Matthews'
def update(self, logits, labels):
labels = labels.asnumpy()
labels = np.reshape(labels, -1)
logits = logits.asnumpy()
logits = np.argmax(logits, axis=1)
self.labels_array = np.concatenate([self.labels_array, labels]).astype(np.bool)
self.logits_array = np.concatenate([self.logits_array, logits]).astype(np.bool)
def get_metrics(self):
if len(self.labels_array) < 2:
return 0.0
tp = np.sum(self.labels_array & self.logits_array)
fp = np.sum(self.labels_array & (~self.logits_array))
fn = np.sum((~self.labels_array) & self.logits_array)
tn = np.sum((~self.labels_array) & (~self.logits_array))
return (tp * tn - fp * fn) / np.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)) * 100.0
This diff is collapsed.
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""config script"""
from easydict import EasyDict as edict
import mindspore.common.dtype as mstype
from .tinybert_model import BertConfig
from .assessment_method import Accuracy, F1, Pearsonr, Matthews
gradient_cfg = edict({
'clip_type': 1,
'clip_value': 1.0
})
task_cfg = edict({
"sst-2": edict({"num_labels": 2, "seq_length": 64, "task_type": "classification", "metrics": Accuracy}),
"qnli": edict({"num_labels": 2, "seq_length": 128, "task_type": "classification", "metrics": Accuracy}),
"mnli": edict({"num_labels": 3, "seq_length": 128, "task_type": "classification", "metrics": Accuracy}),
"cola": edict({"num_labels": 2, "seq_length": 64, "task_type": "classification", "metrics": Matthews}),
"mrpc": edict({"num_labels": 2, "seq_length": 128, "task_type": "classification", "metrics": F1}),
"sts-b": edict({"num_labels": 1, "seq_length": 128, "task_type": "regression", "metrics": Pearsonr}),
"qqp": edict({"num_labels": 2, "seq_length": 128, "task_type": "classification", "metrics": F1}),
"rte": edict({"num_labels": 2, "seq_length": 128, "task_type": "classification", "metrics": Accuracy})
})
train_cfg = edict({
'batch_size': 16,
'loss_scale_value': 2 ** 16,
'scale_factor': 2,
'scale_window': 50,
'optimizer_cfg': edict({
'AdamWeightDecay': edict({
'learning_rate': 5e-5,
'end_learning_rate': 1e-14,
'power': 1.0,
'weight_decay': 1e-4,
'eps': 1e-6,
'decay_filter': lambda x: 'layernorm' not in x.name.lower() and 'bias' not in x.name.lower(),
'warmup_ratio': 0.1
}),
}),
})
eval_cfg = edict({
'batch_size': 32,
})
teacher_net_cfg = BertConfig(
seq_length=128,
vocab_size=30522,
hidden_size=768,
num_hidden_layers=6,
num_attention_heads=12,
intermediate_size=3072,
hidden_act="gelu",
hidden_dropout_prob=0.0,
attention_probs_dropout_prob=0.0,
cls_dropout_prob=0.0,
max_position_embeddings=512,
type_vocab_size=2,
initializer_range=0.02,
use_relative_positions=False,
dtype=mstype.float32,
compute_type=mstype.float32,
do_quant=False
)
student_net_cfg = BertConfig(
seq_length=128,
vocab_size=30522,
hidden_size=768,
num_hidden_layers=6,
num_attention_heads=12,
intermediate_size=3072,
hidden_act="gelu",
hidden_dropout_prob=0.0,
attention_probs_dropout_prob=0.0,
max_position_embeddings=512,
type_vocab_size=2,
initializer_range=0.02,
use_relative_positions=False,
dtype=mstype.float32,
compute_type=mstype.float32,
do_quant=True,
embedding_bits=2,
weight_bits=2,
weight_clip_value=3.0,
cls_dropout_prob=0.0,
activation_init=2.5,
is_lgt_fit=False
)
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""create tinybert dataset"""
from enum import Enum
import mindspore.common.dtype as mstype
import mindspore.dataset.engine as de
import mindspore.dataset.transforms.c_transforms as C
class DataType(Enum):
"""Enumerate supported dataset format"""
TFRECORD = 1
MINDRECORD = 2
def create_dataset(batch_size=32, device_num=1, rank=0, do_shuffle=True, data_dir=None,
data_type='tfrecord', seq_length=128, task_type=mstype.int32, drop_remainder=True):
"""create tinybert dataset"""
if isinstance(data_dir, list):
data_files = data_dir
else:
data_files = [data_dir]
columns_list = ["input_ids", "input_mask", "segment_ids", "label_ids"]
if data_type == 'mindrecord':
ds = de.MindDataset(data_files, columns_list=columns_list, shuffle=do_shuffle, num_shards=device_num,
shard_id=rank)
else:
ds = de.TFRecordDataset(data_files, columns_list=columns_list, shuffle=do_shuffle, num_shards=device_num,
shard_id=rank, shard_equal_rows=(device_num != 1))
if device_num == 1 and do_shuffle is True:
ds = ds.shuffle(10000)
type_cast_op = C.TypeCast(mstype.int32)
slice_op = C.Slice(slice(0, seq_length, 1))
label_type = mstype.int32 if task_type == 'classification' else mstype.float32
ds = ds.map(operations=[type_cast_op, slice_op], input_columns=["segment_ids"])
ds = ds.map(operations=[type_cast_op, slice_op], input_columns=["input_mask"])
ds = ds.map(operations=[type_cast_op, slice_op], input_columns=["input_ids"])
ds = ds.map(operations=[C.TypeCast(label_type), slice_op], input_columns=["label_ids"])
# apply batch operations
ds = ds.batch(batch_size, drop_remainder=drop_remainder)
return ds
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Quantization function."""
from mindspore.common import dtype as mstype
from mindspore.common.parameter import Parameter
from mindspore.ops import operations as P
from mindspore.ops import functional as F
from mindspore.ops import composite as C
from mindspore import nn
class QuantizeWeightCell(nn.Cell):
"""
The ternary fake quant op for weight.
Args:
num_bits (int): The bit number of quantization, supporting 2 to 8 bits. Default: 2.
compute_type (:class:`mindspore.dtype`): Compute type in QuantizeWeightCell. Default: mstype.float32.
clip_value (float): Clips weight to be in [-clip_value, clip_value].
per_channel (bool): Quantization granularity based on layer or on channel. Default: False.
Inputs:
- **weight** (Parameter) - Parameter of shape :math:`(N, C_{in}, H_{in}, W_{in})`.
Outputs:
Parameter of shape :math:`(N, C_{out}, H_{out}, W_{out})`.
"""
def __init__(self, num_bits=8, compute_type=mstype.float32, clip_value=1.0, per_channel=False):
super(QuantizeWeightCell, self).__init__()
self.num_bits = num_bits
self.compute_type = compute_type
self.clip_value = clip_value
self.per_channel = per_channel
self.clamp = C.clip_by_value
self.abs = P.Abs()
self.sum = P.ReduceSum()
self.nelement = F.size
self.div = P.Div()
self.cast = P.Cast()
self.max = P.ReduceMax()
self.min = P.ReduceMin()
self.round = P.Round()
def construct(self, weight):
"""quantize weight cell"""
tensor = self.clamp(weight, -self.clip_value, self.clip_value)
if self.num_bits == 2:
if self.per_channel:
n = self.nelement(tensor[0])
m = self.div(self.sum(self.abs(tensor), 1), n)
thres = 0.7 * m
pos = self.cast(tensor[:] > thres[0], self.compute_type)
neg = self.cast(tensor[:] < -thres[0], self.compute_type)
mask = self.cast(self.abs(tensor)[:] > thres[0], self.compute_type)
alpha = self.reshape(self.sum(self.abs(mask * tensor), 1) / self.sum(mask, 1), (-1, 1))
output = alpha * pos - alpha * neg
else:
n = self.nelement(tensor)
m = self.div(self.sum(self.abs(tensor)), n)
thres = 0.7 * m
pos = self.cast(tensor > thres, self.compute_type)
neg = self.cast(tensor < -thres, self.compute_type)
mask = self.cast(self.abs(tensor) > thres, self.compute_type)
alpha = self.sum(self.abs(mask * self.cast(tensor, self.compute_type))) / self.sum(mask)
output = alpha * pos - alpha * neg
else:
tensor_max = self.cast(self.max(tensor), self.compute_type)
tensor_min = self.cast(self.min(tensor), self.compute_type)
s = (tensor_max - tensor_min) / (2 ** self.cast(self.num_bits, self.compute_type) - 1)
output = self.round(self.div(tensor - tensor_min, s)) * s + tensor_min
return output
class QuantizeWeight:
"""
Quantize weight into specified bit.
Args:
num_bits (int): The bit number of quantization, supporting 2 to 8 bits. Default: 2.
compute_type (:class:`mindspore.dtype`): Compute type in QuantizeWeightCell. Default: mstype.float32.
clip_value (float): Clips weight to be in [-clip_value, clip_value].
per_channel (bool): Quantization granularity based on layer or on channel. Default: False.
Inputs:
- **weight** (Parameter) - Parameter of shape :math:`(N, C_{in}, H_{in}, W_{in})`.
Outputs:
Parameter of shape :math:`(N, C_{out}, H_{out}, W_{out})`.
"""
def __init__(self, num_bits=2, compute_type=mstype.float32, clip_value=1.0, per_channel=False):
self.num_bits = num_bits
self.compute_type = compute_type
self.clip_value = clip_value
self.per_channel = per_channel
self.clamp = C.clip_by_value
self.abs = P.Abs()
self.sum = P.ReduceSum()
self.nelement = F.size
self.div = P.Div()
self.cast = P.Cast()
self.max = P.ReduceMax()
self.min = P.ReduceMin()
self.floor = P.Floor()
def construct(self, weight):
"""quantize weight"""
tensor = self.clamp(weight, -self.clip_value, self.clip_value)
if self.num_bits == 2:
if self.per_channel:
n = self.nelement(tensor[0])
m = self.div(self.sum(self.abs(tensor), 1), n)
thres = 0.7 * m
pos = self.cast(tensor[:] > thres[0], self.compute_type)
neg = self.cast(tensor[:] < -thres[0], self.compute_type)
mask = self.cast(self.abs(tensor)[:] > thres[0], self.compute_type)
alpha = self.reshape(self.sum(self.abs(mask * tensor), 1) / self.sum(mask, 1), (-1, 1))
output = alpha * pos - alpha * neg
else:
n = self.nelement(tensor)
m = self.div(self.sum(self.abs(tensor)), n)
thres = 0.7 * m
pos = self.cast(tensor > thres, self.compute_type)
neg = self.cast(tensor < -thres, self.compute_type)
mask = self.cast(self.abs(tensor) > thres, self.compute_type)
alpha = self.sum(self.abs(mask * tensor)) / self.sum(mask)
output = alpha * pos - alpha * neg
else:
tensor_max = self.max(tensor)
tensor_min = self.min(tensor)
s = (tensor_max - tensor_min) / (2 ** self.num_bits - 1)
output = self.floor(self.div((tensor - tensor_min), s) + 0.5) * s + tensor_min
return output
def convert_network(network, embedding_bits=2, weight_bits=2, clip_value=1.0):
quantize_embedding = QuantizeWeight(num_bits=embedding_bits, clip_value=clip_value)
quantize_weight = QuantizeWeight(num_bits=weight_bits, clip_value=clip_value)
for name, param in network.parameters_and_names():
if 'bert_embedding_lookup' in name and 'min' not in name and 'max' not in name:
quantized_param = quantize_embedding.construct(param)
param.set_data(quantized_param)
elif 'weight' in name and 'dense_1' not in name:
quantized_param = quantize_weight.construct(param)
param.set_data(quantized_param)
def save_params(network):
return {name: Parameter(param, 'saved_params') for name, param in network.parameters_and_names()}
def restore_params(network, params_dict):
for name, param in network.parameters_and_names():
param.set_data(params_dict[name])
This diff is collapsed.
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""ternarybert utils"""
import os
import time
import numpy as np
from mindspore import Tensor
from mindspore.common import dtype as mstype
from mindspore.train.callback import Callback
from mindspore.train.serialization import save_checkpoint
from mindspore.ops import operations as P
from mindspore.nn.learning_rate_schedule import LearningRateSchedule, PolynomialDecayLR, WarmUpLR
from .quant import convert_network, save_params, restore_params
class ModelSaveCkpt(Callback):
"""
Saves checkpoint.
If the loss in NAN or INF terminating training.
Args:
network (Network): The train network for training.
save_ckpt_step (int): The step to save checkpoint.
max_ckpt_num (int): The max checkpoint number.
"""
def __init__(self, network, save_ckpt_step, max_ckpt_num, output_dir, embedding_bits=2, weight_bits=2,
clip_value=1.0):
super(ModelSaveCkpt, self).__init__()
self.count = 0
self.network = network
self.save_ckpt_step = save_ckpt_step
self.max_ckpt_num = max_ckpt_num
self.output_dir = output_dir
if not os.path.exists(output_dir):
os.makedirs(output_dir)
self.embedding_bits = embedding_bits
self.weight_bits = weight_bits
self.clip_value = clip_value
def step_end(self, run_context):
"""step end and save ckpt"""
cb_params = run_context.original_args()
if cb_params.cur_step_num % self.save_ckpt_step == 0:
saved_ckpt_num = cb_params.cur_step_num / self.save_ckpt_step
if saved_ckpt_num > self.max_ckpt_num:
oldest_ckpt_index = saved_ckpt_num - self.max_ckpt_num
path = os.path.join(self.output_dir, "ternary_bert_{}_{}.ckpt".format(int(oldest_ckpt_index),
self.save_ckpt_step))
if os.path.exists(path):
os.remove(path)
params_dict = save_params(self.network)
convert_network(self.network, self.embedding_bits, self.weight_bits, self.clip_value)
save_checkpoint(self.network, os.path.join(self.output_dir,
"ternary_bert_{}_{}.ckpt".format(int(saved_ckpt_num),
self.save_ckpt_step)))
restore_params(self.network, params_dict)
class LossCallBack(Callback):
"""
Monitor the loss in training.
"""
def __init__(self, per_print_times=1):
super(LossCallBack, self).__init__()
if not isinstance(per_print_times, int) or per_print_times < 0:
raise ValueError("print_step must be int and >= 0")
self._per_print_times = per_print_times
def step_end(self, run_context):
"""step end and print loss"""
cb_params = run_context.original_args()
print("epoch: {}, step: {}, outputs are {}".format(cb_params.cur_epoch_num,
cb_params.cur_step_num,
str(cb_params.net_outputs)))
class StepCallBack(Callback):
"""
Monitor the time in training.
"""
def __init__(self):
super(StepCallBack, self).__init__()
self.start_time = 0.0
def step_begin(self, run_context):
self.start_time = time.time()
def step_end(self, run_context):
time_cost = time.time() - self.start_time
cb_params = run_context.original_args()
print("step: {}, second_per_step: {}".format(cb_params.cur_step_num, time_cost))
class EvalCallBack(Callback):
"""Evaluation callback"""
def __init__(self, network, dataset, eval_ckpt_step, save_ckpt_dir, embedding_bits=2, weight_bits=2,
clip_value=1.0, metrics=None):
super(EvalCallBack, self).__init__()
self.network = network
self.global_metrics = 0.0
self.dataset = dataset
self.eval_ckpt_step = eval_ckpt_step
self.save_ckpt_dir = save_ckpt_dir
self.embedding_bits = embedding_bits
self.weight_bits = weight_bits
self.clip_value = clip_value
self.metrics = metrics
if not os.path.exists(save_ckpt_dir):
os.makedirs(save_ckpt_dir)
def step_end(self, run_context):
"""step end and do evaluation"""
cb_params = run_context.original_args()
if cb_params.cur_step_num % self.eval_ckpt_step == 0:
params_dict = save_params(self.network)
convert_network(self.network, self.embedding_bits, self.weight_bits, self.clip_value)
self.network.set_train(False)
callback = self.metrics()
columns_list = ["input_ids", "input_mask", "segment_ids", "label_ids"]
for data in self.dataset:
input_data = []
for i in columns_list:
input_data.append(data[i])
input_ids, input_mask, token_type_id, label_ids = input_data
_, _, logits, _ = self.network(input_ids, token_type_id, input_mask)
callback.update(logits, label_ids)
metrics = callback.get_metrics()
if metrics > self.global_metrics:
self.global_metrics = metrics
eval_model_ckpt_file = os.path.join(self.save_ckpt_dir, 'eval_model.ckpt')
if os.path.exists(eval_model_ckpt_file):
os.remove(eval_model_ckpt_file)
save_checkpoint(self.network, eval_model_ckpt_file)
print('step {}, {} {}, best_{} {}'.format(cb_params.cur_step_num,
callback.name,
metrics,
callback.name,
self.global_metrics))
restore_params(self.network, params_dict)
self.network.set_train(True)
class BertLearningRate(LearningRateSchedule):
"""
Warmup-decay learning rate for Bert network.
"""
def __init__(self, learning_rate, end_learning_rate, warmup_steps, decay_steps, power):
super(BertLearningRate, self).__init__()
self.warmup_flag = False
if warmup_steps > 0:
self.warmup_flag = True
self.warmup_lr = WarmUpLR(learning_rate, warmup_steps)
self.decay_lr = PolynomialDecayLR(learning_rate, end_learning_rate, decay_steps, power)
self.warmup_steps = Tensor(np.array([warmup_steps]).astype(np.float32))
self.greater = P.Greater()
self.one = Tensor(np.array([1.0]).astype(np.float32))
self.cast = P.Cast()
def construct(self, global_step):
decay_lr = self.decay_lr(global_step)
if self.warmup_flag:
is_warmup = self.cast(self.greater(self.warmup_steps, global_step), mstype.float32)
warmup_lr = self.warmup_lr(global_step)
lr = (self.one - is_warmup) * decay_lr + is_warmup * warmup_lr
else:
lr = decay_lr
return lr
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""task distill script"""
import os
import argparse
import ast
from mindspore import context
from mindspore.train.model import Model
from mindspore.nn.optim import AdamWeightDecay
from mindspore import set_seed
from mindspore.train.callback import TimeMonitor
from src.dataset import create_dataset
from src.utils import StepCallBack, ModelSaveCkpt, EvalCallBack, BertLearningRate
from src.config import train_cfg, eval_cfg, teacher_net_cfg, student_net_cfg, task_cfg
from src.cell_wrapper import BertNetworkWithLoss, BertTrainCell
WEIGHTS_NAME = 'eval_model.ckpt'
EVAL_DATA_NAME = 'eval.tf_record'
TRAIN_DATA_NAME = 'train.tf_record'
def parse_args():
"""
parse args
"""
parser = argparse.ArgumentParser(description='ternarybert task distill')
parser.add_argument('--device_target', type=str, default='GPU', choices=['Ascend', 'GPU'],
help='Device where the code will be implemented. (Default: GPU)')
parser.add_argument('--do_eval', type=ast.literal_eval, default=True,
help='Do eval task during training or not. (Default: True)')
parser.add_argument('--epoch_size', type=int, default=3, help='Epoch size for train phase. (Default: 3)')
parser.add_argument('--device_id', type=int, default=0, help='Device id. (Default: 0)')
parser.add_argument('--do_shuffle', type=ast.literal_eval, default=True,
help='Enable shuffle for train dataset. (Default: True)')
parser.add_argument('--enable_data_sink', type=ast.literal_eval, default=True,
help='Enable data sink. (Default: True)')
parser.add_argument('--save_ckpt_step', type=int, default=50,
help='If do_eval is False, the checkpoint will be saved every save_ckpt_step. (Default: 50)')
parser.add_argument('--eval_ckpt_step', type=int, default=50,
help='If do_eval is True, the evaluation will be ran every eval_ckpt_step. (Default: 50)')
parser.add_argument('--max_ckpt_num', type=int, default=10,
help='The number of checkpoints will not be larger than max_ckpt_num. (Default: 10)')
parser.add_argument('--data_sink_steps', type=int, default=100, help='Sink steps for each epoch. (Default: 1)')
parser.add_argument('--teacher_model_dir', type=str, default='', help='The checkpoint directory of teacher model.')
parser.add_argument('--student_model_dir', type=str, default='', help='The checkpoint directory of student model.')
parser.add_argument('--data_dir', type=str, default='', help='Data directory.')
parser.add_argument('--output_dir', type=str, default='./', help='The output checkpoint directory.')
parser.add_argument('--task_name', type=str, default='sts-b', choices=['sts-b', 'qnli', 'mnli'],
help='The name of the task to train. (Default: sts-b)')
parser.add_argument('--dataset_type', type=str, default='tfrecord', choices=['tfrecord', 'mindrecord'],
help='The name of the task to train. (Default: tfrecord)')
parser.add_argument('--seed', type=int, default=1, help='The random seed')
parser.add_argument('--train_batch_size', type=int, default=16, help='Batch size for training')
parser.add_argument('--eval_batch_size', type=int, default=32, help='Eval Batch size in callback')
return parser.parse_args()
def run_task_distill(args_opt):
"""
run task distill
"""
task = task_cfg[args_opt.task_name]
teacher_net_cfg.seq_length = task.seq_length
student_net_cfg.seq_length = task.seq_length
train_cfg.batch_size = args_opt.train_batch_size
eval_cfg.batch_size = args_opt.eval_batch_size
teacher_ckpt = os.path.join(args_opt.teacher_model_dir, args_opt.task_name, WEIGHTS_NAME)
student_ckpt = os.path.join(args_opt.student_model_dir, args_opt.task_name, WEIGHTS_NAME)
train_data_dir = os.path.join(args_opt.data_dir, args_opt.task_name, TRAIN_DATA_NAME)
eval_data_dir = os.path.join(args_opt.data_dir, args_opt.task_name, EVAL_DATA_NAME)
save_ckpt_dir = os.path.join(args_opt.output_dir, args_opt.task_name)
if args_opt.device_target == "Ascend":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
elif args_opt.device_target == "GPU":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target)
else:
raise Exception("Target error, GPU or Ascend is supported.")
rank = 0
device_num = 1
train_dataset = create_dataset(batch_size=train_cfg.batch_size,
device_num=device_num,
rank=rank,
do_shuffle=args_opt.do_shuffle,
data_dir=train_data_dir,
data_type=args_opt.dataset_type,
seq_length=task.seq_length,
task_type=task.task_type,
drop_remainder=True)
dataset_size = train_dataset.get_dataset_size()
print('train dataset size:', dataset_size)
eval_dataset = create_dataset(batch_size=eval_cfg.batch_size,
device_num=device_num,
rank=rank,
do_shuffle=args_opt.do_shuffle,
data_dir=eval_data_dir,
data_type=args_opt.dataset_type,
seq_length=task.seq_length,
task_type=task.task_type,
drop_remainder=False)
print('eval dataset size:', eval_dataset.get_dataset_size())
if args_opt.enable_data_sink:
repeat_count = args_opt.epoch_size * dataset_size // args_opt.data_sink_steps
time_monitor_steps = args_opt.data_sink_steps
else:
repeat_count = args_opt.epoch_size
time_monitor_steps = dataset_size
netwithloss = BertNetworkWithLoss(teacher_config=teacher_net_cfg, teacher_ckpt=teacher_ckpt,
student_config=student_net_cfg, student_ckpt=student_ckpt,
is_training=True, task_type=task.task_type, num_labels=task.num_labels)
params = netwithloss.trainable_params()
optimizer_cfg = train_cfg.optimizer_cfg
lr_schedule = BertLearningRate(learning_rate=optimizer_cfg.AdamWeightDecay.learning_rate,
end_learning_rate=optimizer_cfg.AdamWeightDecay.end_learning_rate,
warmup_steps=int(dataset_size * args_opt.epoch_size *
optimizer_cfg.AdamWeightDecay.warmup_ratio),
decay_steps=int(dataset_size * args_opt.epoch_size),
power=optimizer_cfg.AdamWeightDecay.power)
decay_params = list(filter(optimizer_cfg.AdamWeightDecay.decay_filter, params))
other_params = list(filter(lambda x: not optimizer_cfg.AdamWeightDecay.decay_filter(x), params))
group_params = [{'params': decay_params, 'weight_decay': optimizer_cfg.AdamWeightDecay.weight_decay},
{'params': other_params, 'weight_decay': 0.0},
{'order_params': params}]
optimizer = AdamWeightDecay(group_params, learning_rate=lr_schedule, eps=optimizer_cfg.AdamWeightDecay.eps)
netwithgrads = BertTrainCell(netwithloss, optimizer=optimizer)
if args_opt.do_eval:
eval_dataset = list(eval_dataset.create_dict_iterator())
callback = [TimeMonitor(time_monitor_steps),
EvalCallBack(network=netwithloss.bert,
dataset=eval_dataset,
eval_ckpt_step=args_opt.eval_ckpt_step,
save_ckpt_dir=save_ckpt_dir,
embedding_bits=student_net_cfg.embedding_bits,
weight_bits=student_net_cfg.weight_bits,
clip_value=student_net_cfg.weight_clip_value,
metrics=task.metrics)]
else:
callback = [TimeMonitor(time_monitor_steps), StepCallBack(),
ModelSaveCkpt(network=netwithloss.bert,
save_ckpt_step=args_opt.save_ckpt_step,
max_ckpt_num=args_opt.max_ckpt_num,
output_dir=save_ckpt_dir,
embedding_bits=student_net_cfg.embedding_bits,
weight_bits=student_net_cfg.weight_bits,
clip_value=student_net_cfg.weight_clip_value)]
model = Model(netwithgrads)
model.train(repeat_count, train_dataset, callbacks=callback,
dataset_sink_mode=args_opt.enable_data_sink,
sink_size=args_opt.data_sink_steps)
if __name__ == '__main__':
args = parse_args()
set_seed(args.seed)
run_task_distill(args)
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment