# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. # # 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. import sys import os import numpy as np from sklearn.metrics import f1_score import paddle as P from paddle.nn import functional as F import propeller.paddle as propeller from ernie.tokenizing_ernie import ErnieTokenizer from ernie.modeling_ernie import ErnieModelForSequenceClassification from demo.utils import create_if_not_exists, get_warmup_and_linear_decay # 本例子采用chnsenticorp中文情感识别任务作为示范；并且事先通过数据增强扩充了蒸馏所需的无监督数据 # # 下载数据；并存放在 ./chnsenticorp-data/ # 数据分为3列：原文；空格切词；情感标签 # 其中第一列为ERNIE的输入；第二列为BoW词袋模型的输入 # 事先统计好的BoW 词典在 ./chnsenticorp-data/vocab.bow.txt # 定义finetune teacher模型所需要的超参数 DATA_DIR = './chnsenticorp-data/' SEQLEN = 256 BATCH = 32 EPOCH = 10 LR = 5e-5 tokenizer = ErnieTokenizer.from_pretrained('ernie-1.0') student_vocab = { i.strip(): l for l, i in enumerate( open( os.path.join(DATA_DIR, 'vocab.bow.txt'), encoding='utf8') .readlines()) } def space_tokenizer(i): return i.decode('utf8').split() feature_column = propeller.data.FeatureColumns([ propeller.data.TextColumn( 'seg_a', unk_id=tokenizer.unk_id, vocab_dict=tokenizer.vocab, tokenizer=tokenizer.tokenize), propeller.data.TextColumn( 'seg_a_student', unk_id=student_vocab['[UNK]'], vocab_dict=student_vocab, tokenizer=space_tokenizer), propeller.data.LabelColumn( 'label', vocab_dict={ b"0": 0, b"1": 1, }), ]) def map_fn(seg_a, seg_a_student, label): seg_a, _ = tokenizer.truncate(seg_a, [], seqlen=SEQLEN) sentence, segments = tokenizer.build_for_ernie(seg_a) return seg_a_student, sentence, segments, label train_ds = feature_column.build_dataset('train', data_dir=os.path.join(DATA_DIR, 'train/'), shuffle=True, repeat=False, use_gz=False) \ .map(map_fn) \ .padded_batch(BATCH) train_ds_unlabel = feature_column.build_dataset('train-da', data_dir=os.path.join(DATA_DIR, 'train-data-augmented/'), shuffle=True, repeat=False, use_gz=False) \ .map(map_fn) \ .padded_batch(BATCH) dev_ds = feature_column.build_dataset('dev', data_dir=os.path.join(DATA_DIR, 'dev/'), shuffle=False, repeat=False, use_gz=False) \ .map(map_fn) \ .padded_batch(BATCH,) shapes = ([-1, SEQLEN], [-1, SEQLEN], [-1, SEQLEN], [-1]) types = ('int64', 'int64', 'int64', 'int64') train_ds.data_shapes = shapes train_ds.data_types = types train_ds_unlabel.data_shapes = shapes train_ds_unlabel.data_types = types dev_ds.data_shapes = shapes dev_ds.data_types = types place = P.CUDAPlace(0) def evaluate_teacher(model, dataset): all_pred, all_label = [], [] with P.no_grad(): model.eval() for step, (ids_student, ids, _, labels) in enumerate( P.io.DataLoader( dataset, places=place, batch_size=None)): _, logits = model(ids) pred = logits.argmax(-1) all_pred.extend(pred.numpy()) all_label.extend(labels.numpy()) f1 = f1_score(all_label, all_pred, average='macro') model.train() return f1 teacher_model = ErnieModelForSequenceClassification.from_pretrained( 'ernie-1.0', num_labels=2) teacher_model.train() if not os.path.exists('./teacher_model.bin'): g_clip = P.nn.ClipGradByGlobalNorm(1.0) #experimental lr_scheduler = P.optimizer.lr.LambdaDecay( LR, get_warmup_and_linear_decay(9600 * EPOCH / BATCH, 9600 * EPOCH * 0.1 / BATCH)) opt = P.optimizer.AdamW( lr_scheduler, parameters=teacher_model.parameters(), weight_decay=0.01, grad_clip=g_clip) for epoch in range(EPOCH): for step, (ids_student, ids, sids, labels) in enumerate( P.io.DataLoader( train_ds, places=place, batch_size=None)): loss, logits = teacher_model(ids, labels=labels) loss.backward() opt.step() lr_scheduler.step() teacher_model.clear_gradients() if step % 10 == 0: _lr = lr_scheduler.get_lr() _l = loss.numpy() msg = '[step-%d] train loss %.5f lr %.3e' % (step, _l, _lr) print(msg) if step % 100 == 0: f1 = evaluate_teacher(teacher_model, dev_ds) print('teacher f1: %.5f' % f1) P.save(teacher_model.state_dict(), './teacher_model.bin') else: state_dict = P.load('./teacher_model.bin') teacher_model.set_state_dict(state_dict) f1 = evaluate_teacher(teacher_model, dev_ds) print('teacher f1: %.5f' % f1) # 定义finetune student 模型所需要的超参数 SEQLEN = 256 BATCH = 32 EPOCH = 10 LR = 1e-4 def evaluate_student(model, dataset): all_pred, all_label = [], [] with P.no_grad(): model.eval() for step, (ids_student, ids, _, labels) in enumerate( P.io.DataLoader( dataset, places=place, batch_size=None)): _, logits = model(ids_student) pred = logits.argmax(-1) all_pred.extend(pred.numpy()) all_label.extend(labels.numpy()) f1 = f1_score(all_label, all_pred, average='macro') model.train() return f1 class BOW(P.nn.Layer): def __init__(self): super().__init__() self.emb = P.nn.Embedding(len(student_vocab), 128, padding_idx=0) self.fc = P.nn.Linear(128, 2) def forward(self, ids, labels=None): embbed = self.emb(ids) pad_mask = (ids != 0).cast('float32').unsqueeze(-1) embbed = (embbed * pad_mask).sum(1) embbed = F.softsign(embbed) logits = self.fc(embbed) if labels is not None: if len(labels.shape) == 1: labels = labels.reshape([-1, 1]) loss = F.cross_entropy(logits, labels).mean() else: loss = None return loss, logits class CNN(P.nn.Layer): def __init__(self): super().__init__() self.emb = P.nn.Embedding(30002, 128, padding_idx=0) self.cnn = P.nn.Conv2D(128, 128, (1, 3), padding=(0, 1), act='relu') self.pool = P.nn.Pool2D((1, 3), pool_padding=(0, 1)) self.fc = P.nn.Linear(128, 2) def forward(self, ids, labels=None): embbed = self.emb(ids) #d_batch, d_seqlen = ids.shape hidden = embbed hidden = hidden.transpose([0, 2, 1]).unsqueeze(2) #change to NCWH hidden = self.cnn(hidden) hidden = self.pool(hidden).squeeze(2).transpose([0, 2, 1]) pad_mask = (ids != 0).cast('float32').unsqueeze(-1) hidden = P.nn.funcional.softsign(L(hidden * pad_mask).sum(1)) logits = self.fc(hidden) if labels is not None: if len(labels.shape) == 1: labels = labels.reshape([-1, 1]) loss = F.cross_entropy(logits, labels).mean() else: loss = None return loss, logits def KL(pred, target): pred = F.log_softmax(pred) target = F.softmax(target) loss = F.kl_div(pred, target) return loss teacher_model.eval() model = BOW() g_clip = P.nn.ClipGradByGlobalNorm(1.0) #experimental lr_scheduler = P.optimizer.lr.LambdaDecay( LR, get_warmup_and_linear_decay(9600 * EPOCH / BATCH, 9600 * EPOCH * 0.1 / BATCH)) opt = P.optimizer.AdamW( lr_scheduler, parameters=model.parameters(), weight_decay=0.01, grad_clip=g_clip) model.train() for epoch in range(EPOCH - 1): for step, ( ids_student, ids, sids, label ) in enumerate(P.io.DataLoader( train_ds, places=place, batch_size=None)): with P.no_grad(): _, logits_t = teacher_model(ids, sids) # teacher 模型输出logits _, logits_s = model(ids_student) # student 模型输出logits loss_ce, _ = model(ids_student, labels=label) loss_kd = KL(logits_s, logits_t.detach()) # 由KL divergence度量两个分布的距离 loss = loss_ce + loss_kd loss.backward() opt.step() lr_scheduler.step() model.clear_gradients() if step % 10 == 0: _lr = lr_scheduler.get_lr() _l = loss.numpy() msg = '[step-%d] train loss %.5f lr %.3e' % (step, _l, _lr) print(msg) f1 = evaluate_student(model, dev_ds) print('student f1 %.5f' % f1) # 最后再加一轮hard label训练巩固结果 for step, ( ids_student, ids, sids, label ) in enumerate(P.io.DataLoader( train_ds, places=place, batch_size=None)): loss, _ = model(ids_student, labels=label) loss.backward() opt.step() model.clear_gradients() if step % 10 == 0: _lr = lr_scheduler.get_lr() _l = loss.numpy() msg = '[step-%d] train loss %.5f lr %.3e' % (step, _l, _lr) print(msg) f1 = evaluate_student(model, dev_ds) print('final f1 %.5f' % f1)