以下是本例的简要目录结构及说明:
.
├── student_lstm_small.py # 小模型结构以及对小模型单独训练的脚本
├── teacher_train_glue.py # 训练教师模型的脚本
├── teacher_student_distill.py # 用教师模型BERT蒸馏学生模型的蒸馏脚本
├── data.py # 定义了dataloader等数据读取接口
├── utils.py # 定义了将样本转成id的转换接口
├── args.py # 参数配置脚本
└── README.md # 文档,本文件
wget https://paddlenlp.bj.bcebos.com/data/senta_word_dict.txt
export TASK_NAME=ChnSentiCorp
python -u ./teacher_train_glue.py \
--model_type ernie-gram \
--model_name_or_path ernie-gram-zh \
--task_name $TASK_NAME \
--max_seq_length 128 \
--batch_size 32 \
--learning_rate 5e-5 \
--num_train_epochs 30 \
--logging_steps 10 \
--save_steps 100 \
--output_dir ./teacher/$TASK_NAME/ \
--device gpu # or cpu
这一步是将教师模型ernie-gram-zh的知识蒸馏到基于BiLSTM的学生模型中,可以运行下面的命令分别基于ChnSentiCorp对基于BiLSTM的学生模型进行蒸馏。
CUDA_VISIBLE_DEVICES=0
python teacher_student_distill.py \
--task_name chnsenticorp \
--vocab_size 1256608 \
--max_epoch 6 \
--lr 1.0 \
--dropout_prob 0.1 \
--batch_size 64 \
--model_name ernie-gram-zh \
--teacher_dir ./teacher/ChnSentiCorp/chnsenticorp_ft_model_9000.pdparams \
--vocab_path senta_word_dict.txt \
--output_dir distilled_models \
--save_steps 10000 \
各参数的具体说明请参阅 args.py
,注意在训练不同任务时,需要调整对应的超参数。
本蒸馏实验基于GLUE的中文情感分类ChnSentiCorp数据集。实验效果均使用每个数据集的验证集(dev)进行评价,评价指标是准确率(acc),其中QQP中包含f1值。利用基于ernie-gram-zh的教师模型去蒸馏基于Bi-LSTM的学生模型,对比Bi-LSTM小模型单独训练,在ChnSentiCorp(中文情感分类)任务上分别有1.4%的提升。
Tang R, Lu Y, Liu L, Mou L, Vechtomova O, Lin J. Distilling Task-Specific Knowledge from BERT into Simple Neural Networks[J]. arXiv preprint arXiv:1903.12136, 2019.