class MTDNNModel(object):
def __init__(self, opt, state_dict=None, num_train_step=-1):
self.config = opt
self.updates = state_dict[
'updates'] if state_dict and 'updates' in state_dict else 0
self.train_loss = AverageMeter()
self.network = SANBertNetwork(opt)
# pdb.set_trace()
if state_dict:
new_state = set(self.network.state_dict().keys())
# change to a safer approach
old_keys = [k for k in state_dict['state'].keys()]
for k in old_keys:
if k not in new_state:
print('deleting state:', k)
del state_dict['state'][k]
for k, v in list(self.network.state_dict().items()):
if k not in state_dict['state']:
print('adding missing state:', k)
state_dict['state'][k] = v
# pdb.set_trace()
self.network.load_state_dict(state_dict['state'])
self.mnetwork = nn.DataParallel(
self.network) if opt['multi_gpu_on'] else self.network
self.total_param = sum([
p.nelement() for p in self.network.parameters() if p.requires_grad
])
no_decay = [
'bias', 'gamma', 'beta', 'LayerNorm.bias', 'LayerNorm.weight'
]
optimizer_parameters = [{
'params': [
p for n, p in self.network.named_parameters()
if n not in no_decay
],
'weight_decay_rate':
0.01
}, {
'params':
[p for n, p in self.network.named_parameters() if n in no_decay],
'weight_decay_rate':
0.0
}]
# note that adamax are modified based on the BERT code
if opt['optimizer'] == 'sgd':
self.optimizer = optim.SGD(optimizer_parameters,
opt['learning_rate'],
weight_decay=opt['weight_decay'])
elif opt['optimizer'] == 'adamax':
self.optimizer = Adamax(optimizer_parameters,
opt['learning_rate'],
warmup=opt['warmup'],
t_total=num_train_step,
max_grad_norm=opt['grad_clipping'],
schedule=opt['warmup_schedule'])
if opt.get('have_lr_scheduler', False):
opt['have_lr_scheduler'] = False
elif opt['optimizer'] == 'adadelta':
self.optimizer = optim.Adadelta(optimizer_parameters,
opt['learning_rate'],
rho=0.95)
elif opt['optimizer'] == 'adam':
self.optimizer = Adam(optimizer_parameters,
lr=opt['learning_rate'],
warmup=opt['warmup'],
t_total=num_train_step,
max_grad_norm=opt['grad_clipping'],
schedule=opt['warmup_schedule'])
if opt.get('have_lr_scheduler', False):
opt['have_lr_scheduler'] = False
else:
raise RuntimeError('Unsupported optimizer: %s' % opt['optimizer'])
if state_dict and 'optimizer' in state_dict:
self.optimizer.load_state_dict(state_dict['optimizer'])
if opt.get('have_lr_scheduler', False):
if opt.get('scheduler_type', 'rop') == 'rop':
self.scheduler = ReduceLROnPlateau(self.optimizer,
mode='max',
factor=opt['lr_gamma'],
patience=3)
elif opt.get('scheduler_type', 'rop') == 'exp':
self.scheduler = ExponentialLR(self.optimizer,
gamma=opt.get('lr_gamma', 0.95))
else:
milestones = [
int(step)
for step in opt.get('multi_step_lr', '10,20,30').split(',')
]
self.scheduler = MultiStepLR(self.optimizer,
milestones=milestones,
gamma=opt.get('lr_gamma'))
else:
self.scheduler = None
self.ema = None
if opt['ema_opt'] > 0:
self.ema = EMA(self.config['ema_gamma'], self.network)
self.para_swapped = False
def setup_ema(self):
if self.config['ema_opt']:
self.ema.setup()
def update_ema(self):
if self.config['ema_opt']:
self.ema.update()
def eval(self):
if self.config['ema_opt']:
self.ema.swap_parameters()
self.para_swapped = True
def train(self):
if self.para_swapped:
self.ema.swap_parameters()
self.para_swapped = False
def update(self, batch_meta, batch_data):
self.network.train()
labels = batch_data[batch_meta['label']]
# print('data size:',batch_data[batch_meta['token_id']].size())
if batch_meta['pairwise']:
labels = labels.contiguous().view(-1,
batch_meta['pairwise_size'])[:,
0]
if self.config['cuda']:
y = Variable(labels.cuda(async=True), requires_grad=False)
else:
y = Variable(labels, requires_grad=False)
task_id = batch_meta['task_id']
task_type = batch_meta['task_type']
inputs = batch_data[:batch_meta['input_len']]
if len(inputs) == 3:
inputs.append(None)
inputs.append(None)
inputs.append(task_id)
# pdb.set_trace()
logits = self.mnetwork(*inputs)
if batch_meta['pairwise']:
logits = logits.view(-1, batch_meta['pairwise_size'])
# pdb.set_trace()
if task_type > 0:
if self.config['answer_relu']:
logits = F.relu(logits)
loss = F.mse_loss(logits.squeeze(1), y)
else:
loss = F.cross_entropy(logits, y)
if self.config['mediqa_pairloss'] is not None and batch_meta[
'dataset_name'] in mediqa_name_list:
# print(logits)
# print(batch_data[batch_meta['rank_label']].size())
# input('ha')
logits = logits.squeeze().view(-1, 2)
# print(batch_data[batch_meta['rank_label']])
rank_y = batch_data[batch_meta['rank_label']].view(-1, 2)
# print(rank_y)
if self.config['mediqa_pairloss'] == 'hinge':
# print(logits)
first_logit, second_logit = logits.split(1, dim=1)
# print(first_logit,second_logit)
# pdb.set_trace()
rank_y = (2 * rank_y - 1).to(torch.float32)
rank_y = rank_y[:, 0]
pairwise_loss = F.margin_ranking_loss(
first_logit.squeeze(1),
second_logit.squeeze(1),
rank_y,
margin=self.config['hinge_lambda'])
else:
# pdb.set_trace()
pairwise_loss = F.cross_entropy(logits, rank_y[:, 1])
# print('pairwise_loss:',pairwise_loss,'mse loss:',loss)
loss += pairwise_loss
self.train_loss.update(loss.item(), logits.size(0))
self.optimizer.zero_grad()
loss.backward()
if self.config['global_grad_clipping'] > 0:
torch.nn.utils.clip_grad_norm_(self.network.parameters(),
self.config['global_grad_clipping'])
self.optimizer.step()
self.updates += 1
self.update_ema()
def predict(self, batch_meta, batch_data):
self.network.eval()
task_id = batch_meta['task_id']
task_type = batch_meta['task_type']
inputs = batch_data[:batch_meta['input_len']]
if len(inputs) == 3:
inputs.append(None)
inputs.append(None)
inputs.append(task_id)
score = self.mnetwork(*inputs)
gold_label = batch_meta['label']
if batch_meta['pairwise']:
score = score.contiguous().view(-1, batch_meta['pairwise_size'])
if task_type < 1:
score = F.softmax(score, dim=1)
score = score.data.cpu()
score = score.numpy()
predict = np.zeros(score.shape, dtype=int)
if task_type < 1:
positive = np.argmax(score, axis=1)
for idx, pos in enumerate(positive):
predict[idx, pos] = 1
predict = predict.reshape(-1).tolist()
score = score.reshape(-1).tolist()
return score, predict, batch_meta['true_label']
else:
if task_type < 1:
score = F.softmax(score, dim=1)
# pdb.set_trace()
score = score.data.cpu()
score = score.numpy()
if task_type < 1:
predict = np.argmax(score, axis=1).tolist()
else:
predict = np.greater(
score,
2.0 + self.config['mediqa_score_offset']).astype(int)
gold_label = np.greater(
batch_meta['label'],
2.00001 + self.config['mediqa_score_offset']).astype(int)
predict = predict.reshape(-1).tolist()
gold_label = gold_label.reshape(-1).tolist()
# print('predict:',predict,score)
score = score.reshape(-1).tolist()
return score, predict, gold_label
def save(self, filename):
network_state = dict([(k, v.cpu())
for k, v in self.network.state_dict().items()])
ema_state = dict([
(k, v.cpu()) for k, v in self.ema.model.state_dict().items()
]) if self.ema is not None else dict()
params = {
'state': network_state,
'optimizer': self.optimizer.state_dict(),
'ema': ema_state,
'config': self.config,
}
torch.save(params, filename)
logger.info('model saved to {}'.format(filename))
def cuda(self):
self.network.cuda()
if self.config['ema_opt']:
self.ema.cuda()
def main(*_, **kwargs):
use_cuda = torch.cuda.is_available() and kwargs["device"] >= 0
device = torch.device("cuda:" +
str(kwargs["device"]) if use_cuda else "cpu")
if use_cuda:
torch.cuda.set_device(device)
kwargs["use_cuda"] = use_cuda
neptune.create_experiment(
name="bert-span-parser",
upload_source_files=[],
params={
k: str(v) if isinstance(v, bool) else v
for k, v in kwargs.items()
},
)
logger.info("Settings: {}", json.dumps(kwargs,
indent=2,
ensure_ascii=False))
# For reproducibility
os.environ["PYTHONHASHSEED"] = str(kwargs["seed"])
random.seed(kwargs["seed"])
np.random.seed(kwargs["seed"])
torch.manual_seed(kwargs["seed"])
torch.cuda.manual_seed_all(kwargs["seed"])
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Prepare and load data
tokenizer = BertTokenizer.from_pretrained(kwargs["bert_model"],
do_lower_case=False)
logger.info("Loading data...")
train_treebank = load_trees(kwargs["train_file"])
dev_treebank = load_trees(kwargs["dev_file"])
test_treebank = load_trees(kwargs["test_file"])
logger.info(
"Loaded {:,} train, {:,} dev, and {:,} test examples!",
len(train_treebank),
len(dev_treebank),
len(test_treebank),
)
logger.info("Preprocessing data...")
train_parse = [tree.convert() for tree in train_treebank]
train_sentences = [[(leaf.tag, leaf.word) for leaf in tree.leaves()]
for tree in train_parse]
dev_sentences = [[(leaf.tag, leaf.word) for leaf in tree.leaves()]
for tree in dev_treebank]
test_sentences = [[(leaf.tag, leaf.word) for leaf in tree.leaves()]
for tree in test_treebank]
logger.info("Data preprocessed!")
logger.info("Preparing data for training...")
tags = []
labels = []
for tree in train_parse:
nodes = [tree]
while nodes:
node = nodes.pop()
if isinstance(node, InternalParseNode):
labels.append(node.label)
nodes.extend(reversed(node.children))
else:
tags.append(node.tag)
tag_encoder = LabelEncoder()
tag_encoder.fit(tags, reserved_labels=["[PAD]", "[UNK]"])
label_encoder = LabelEncoder()
label_encoder.fit(labels, reserved_labels=[()])
logger.info("Data prepared!")
# Settings
num_train_optimization_steps = kwargs["num_epochs"] * (
(len(train_parse) - 1) // kwargs["batch_size"] + 1)
kwargs["batch_size"] //= kwargs["gradient_accumulation_steps"]
logger.info("Creating dataloaders for training...")
train_dataloader, train_features = create_dataloader(
sentences=train_sentences,
batch_size=kwargs["batch_size"],
tag_encoder=tag_encoder,
tokenizer=tokenizer,
is_eval=False,
)
dev_dataloader, dev_features = create_dataloader(
sentences=dev_sentences,
batch_size=kwargs["batch_size"],
tag_encoder=tag_encoder,
tokenizer=tokenizer,
is_eval=True,
)
test_dataloader, test_features = create_dataloader(
sentences=test_sentences,
batch_size=kwargs["batch_size"],
tag_encoder=tag_encoder,
tokenizer=tokenizer,
is_eval=True,
)
logger.info("Dataloaders created!")
# Initialize model
model = ChartParser.from_pretrained(
kwargs["bert_model"],
tag_encoder=tag_encoder,
label_encoder=label_encoder,
lstm_layers=kwargs["lstm_layers"],
lstm_dim=kwargs["lstm_dim"],
tag_embedding_dim=kwargs["tag_embedding_dim"],
label_hidden_dim=kwargs["label_hidden_dim"],
dropout_prob=kwargs["dropout_prob"],
)
model.to(device)
# Prepare optimizer
param_optimizers = list(model.named_parameters())
if kwargs["freeze_bert"]:
for p in model.bert.parameters():
p.requires_grad = False
param_optimizers = [(n, p) for n, p in param_optimizers
if p.requires_grad]
# Hack to remove pooler, which is not used thus it produce None grad that break apex
param_optimizers = [n for n in param_optimizers if "pooler" not in n[0]]
no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
{
"params": [
p for n, p in param_optimizers
if not any(nd in n for nd in no_decay)
],
"weight_decay":
0.01,
},
{
"params": [
p for n, p in param_optimizers
if any(nd in n for nd in no_decay)
],
"weight_decay":
0.0,
},
]
optimizer = BertAdam(
optimizer_grouped_parameters,
lr=kwargs["learning_rate"],
warmup=kwargs["warmup_proportion"],
t_total=num_train_optimization_steps,
)
if kwargs["fp16"]:
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
pretrained_model_file = os.path.join(kwargs["output_dir"], MODEL_FILENAME)
if kwargs["do_eval"]:
assert os.path.isfile(
pretrained_model_file), "Pretrained model file does not exist!"
logger.info("Loading pretrained model from {}", pretrained_model_file)
# Load model from file
params = torch.load(pretrained_model_file, map_location=device)
model.load_state_dict(params["model"])
logger.info(
"Loaded pretrained model (Epoch: {:,}, Fscore: {:.2f})",
params["epoch"],
params["fscore"],
)
eval_score = eval(
model=model,
eval_dataloader=test_dataloader,
eval_features=test_features,
eval_trees=test_treebank,
eval_sentences=test_sentences,
tag_encoder=tag_encoder,
device=device,
)
neptune.send_metric("test_eval_precision", eval_score.precision())
neptune.send_metric("test_eval_recall", eval_score.recall())
neptune.send_metric("test_eval_fscore", eval_score.fscore())
tqdm.write("Evaluation score: {}".format(str(eval_score)))
else:
# Training phase
global_steps = 0
start_epoch = 0
best_dev_fscore = 0
if kwargs["preload"] or kwargs["resume"]:
assert os.path.isfile(
pretrained_model_file), "Pretrained model file does not exist!"
logger.info("Resuming model from {}", pretrained_model_file)
# Load model from file
params = torch.load(pretrained_model_file, map_location=device)
model.load_state_dict(params["model"])
if kwargs["resume"]:
optimizer.load_state_dict(params["optimizer"])
torch.cuda.set_rng_state_all([
state.cpu()
for state in params["torch_cuda_random_state_all"]
])
torch.set_rng_state(params["torch_random_state"].cpu())
np.random.set_state(params["np_random_state"])
random.setstate(params["random_state"])
global_steps = params["global_steps"]
start_epoch = params["epoch"] + 1
best_dev_fscore = params["fscore"]
else:
assert not os.path.isfile(
pretrained_model_file
), "Please remove or move the pretrained model file to another place!"
for epoch in trange(start_epoch, kwargs["num_epochs"], desc="Epoch"):
model.train()
train_loss = 0
num_train_steps = 0
for step, (indices, *_) in enumerate(
tqdm(train_dataloader, desc="Iteration")):
ids, attention_masks, tags, sections, trees, sentences = prepare_batch_input(
indices=indices,
features=train_features,
trees=train_parse,
sentences=train_sentences,
tag_encoder=tag_encoder,
device=device,
)
loss = model(
ids=ids,
attention_masks=attention_masks,
tags=tags,
sections=sections,
sentences=sentences,
gold_trees=trees,
)
if kwargs["gradient_accumulation_steps"] > 1:
loss /= kwargs["gradient_accumulation_steps"]
if kwargs["fp16"]:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
train_loss += loss.item()
num_train_steps += 1
if (step + 1) % kwargs["gradient_accumulation_steps"] == 0:
optimizer.step()
optimizer.zero_grad()
global_steps += 1
# Write logs
neptune.send_metric("train_loss", epoch,
train_loss / num_train_steps)
neptune.send_metric("global_steps", epoch, global_steps)
tqdm.write(
"Epoch: {:,} - Train loss: {:.4f} - Global steps: {:,}".format(
epoch, train_loss / num_train_steps, global_steps))
# Evaluate
eval_score = eval(
model=model,
eval_dataloader=dev_dataloader,
eval_features=dev_features,
eval_trees=dev_treebank,
eval_sentences=dev_sentences,
tag_encoder=tag_encoder,
device=device,
)
neptune.send_metric("eval_precision", epoch,
eval_score.precision())
neptune.send_metric("eval_recall", epoch, eval_score.recall())
neptune.send_metric("eval_fscore", epoch, eval_score.fscore())
tqdm.write("Epoch: {:,} - Evaluation score: {}".format(
epoch, str(eval_score)))
# Save best model
if eval_score.fscore() > best_dev_fscore:
best_dev_fscore = eval_score.fscore()
tqdm.write("** Saving model...")
os.makedirs(kwargs["output_dir"], exist_ok=True)
torch.save(
{
"epoch":
epoch,
"global_steps":
global_steps,
"fscore":
best_dev_fscore,
"random_state":
random.getstate(),
"np_random_state":
np.random.get_state(),
"torch_random_state":
torch.get_rng_state(),
"torch_cuda_random_state_all":
torch.cuda.get_rng_state_all(),
"optimizer":
optimizer.state_dict(),
"model": (model.module if hasattr(model, "module") else
model).state_dict(),
},
pretrained_model_file,
)
tqdm.write(
"** Best evaluation fscore: {:.2f}".format(best_dev_fscore))
for epoch in range(args.epoch_num):
## Train
model.train()
t = trange(args.steps_per_epoch,
desc='Epoch {} -Train'.format(epoch))
loss_avg = utils.RunningAverage()
train_iters = [iter(tmp) for tmp in train_bls
] # to use next and reset the iterator
for i in t:
task_id = train_task_ids[i]
batch_data = next(train_iters[task_id])
batch_data = tuple(tmp.to(args.device) for tmp in batch_data)
loss = model(batch_data, task_id, True)
loss.backward()
optimizer.step()
optimizer.zero_grad()
loss_avg.update(loss.item())
t.set_postfix(loss='{:5.4f}'.format(loss.item()),
avg_loss='{:5.4f}'.format(loss_avg()))
acc = eval(model, ner_dev_data, dev_data, dev_bl, graph,
entity_linking, args)
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict()
},
is_best=acc > best_acc,
checkpoint=args.model_dir)
best_acc = max(best_acc, acc)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--device",
default=None,
type=str,
required=True,
help="The GPU device you will run on.")
parser.add_argument(
"--features_file",
default=None,
type=str,
required=True,
help=
"The train features file. Should contain the .csv files (after tokenized) for the task."
"Format: example_id,input_ids,input_mask,segment_ids,label\n")
parser.add_argument(
"--teacher_model",
default=None,
type=str,
help=
"The teacher model dir. Should contain the config/vocab/checkpoint file."
)
parser.add_argument(
"--general_student_model",
default=None,
type=str,
required=True,
help="The student model (after general distillation) dir. "
"Should contain the config/vocab/checkpoint file.")
parser.add_argument(
"--output_student_dir",
default=None,
type=str,
required=True,
help=
"The output directory for the task-specific distilled student models.")
parser.add_argument("--cache_file_dir",
default='./cache',
type=str,
required=True,
help="The directory where cache the features.")
parser.add_argument(
"--distill_model",
default='simplified',
type=str,
help="The distill model type, choose in 'standard' and 'simplified'.")
parser.add_argument(
"--max_seq_length",
default=256,
type=int,
help=
"The maximum total input sequence length after WordPiece tokenization."
)
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=64,
type=int,
help="Total batch size for training.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument('--weight_decay',
'--wd',
default=1e-2,
type=float,
metavar='W',
help='weight decay')
parser.add_argument("--num_train_epochs",
default=2,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--alpha",
default=0.5,
type=float,
help="The weight of soft loss in standard kd method."
"Only use when '--distill_model' is set as 'standard'.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument('--seed',
type=int,
default=42,
help="Random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument(
'--train_loss_step',
type=int,
default=1000,
help="How many train step to record a training loss. ")
parser.add_argument('--save_model_step',
type=int,
default=3000,
help="How many train step to save a student model.")
parser.add_argument('--temperature',
type=float,
default=1.,
help="The temperature in soft loss.")
parser.add_argument(
'--fp16',
action='store_true',
help=
"Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit."
)
parser.add_argument(
'--fp16_opt_level',
type=str,
default='O1',
help=
"For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
"See details at https://nvidia.github.io/apex/amp.html")
args = parser.parse_args()
logger.info('The args: {}'.format(args))
# Prepare device
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
device = torch.device(
"cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
logger.info("device: {} n_gpu: {}".format(device, n_gpu))
# Prepare seed
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
# Prepare task settings
if os.path.exists(args.output_student_dir) and os.listdir(
args.output_student_dir):
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_student_dir))
if not os.path.exists(args.output_student_dir):
os.makedirs(args.output_student_dir)
if not os.path.exists(args.cache_file_dir):
os.makedirs(args.cache_file_dir)
# For save vocab file for all output models.
tokenizer = BertTokenizer.from_pretrained(args.general_student_model,
do_lower_case=args.do_lower_case)
# Model
teacher_model = TinyBertForSequenceClassification.from_pretrained(
args.teacher_model, num_labels=2)
if args.fp16:
teacher_model.half()
teacher_model.to(device)
student_model = TinyBertForSequenceClassification.from_pretrained(
args.general_student_model, num_labels=2)
student_model.to(device)
# Train Config
num_examples, train_dataloader = distill_dataloader(
args, RandomSampler, batch_size=args.train_batch_size)
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
num_train_optimization_steps = int(
num_examples / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
logger.info("***** Running Distilling *****")
logger.info(" Num examples = %d", num_examples)
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
# Prepare optimizer
param_optimizer = list(student_model.named_parameters())
size = 0
for n, p in student_model.named_parameters():
logger.info('n: {}'.format(n))
size += p.nelement()
logger.info('Total parameters of student_model: {}'.format(size))
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
schedule = 'warmup_linear'
optimizer = BertAdam(optimizer_grouped_parameters,
schedule=schedule,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
student_model, optimizer = amp.initialize(
student_model, optimizer, opt_level=args.fp16_opt_level)
logger.info('FP16 is activated, use amp')
else:
logger.info('FP16 is not activated, only use BertAdam')
if n_gpu > 1:
student_model = torch.nn.DataParallel(student_model)
teacher_model = torch.nn.DataParallel(teacher_model)
# Prepare loss functions
loss_mse = MSELoss()
def soft_cross_entropy(predicts, targets):
student_likelihood = torch.nn.functional.log_softmax(predicts, dim=-1)
targets_prob = torch.nn.functional.softmax(targets, dim=-1)
return (-targets_prob * student_likelihood).mean()
# Train
global_step = 0
output_loss_file = os.path.join(args.output_student_dir, "train_loss.txt")
tr_loss = 0.
tr_att_loss = 0.
tr_rep_loss = 0.
tr_cls_loss = 0.
for epoch in trange(int(args.num_train_epochs), desc="Epoch"):
student_model.train()
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration", ascii=True)):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
if input_ids.size()[0] != args.train_batch_size:
continue
student_logits, student_atts, student_reps = student_model(
input_ids, segment_ids, input_mask, is_student=True)
with torch.no_grad():
teacher_logits, teacher_atts, teacher_reps = teacher_model(
input_ids, segment_ids, input_mask)
soft_loss = soft_cross_entropy(student_logits / args.temperature,
teacher_logits / args.temperature)
hard_loss = torch.nn.functional.cross_entropy(student_logits,
label_ids,
reduction='mean')
if args.distill_model == 'standard':
cls_loss = args.alpha * soft_loss + (1 -
args.alpha) * hard_loss
tr_cls_loss += cls_loss.item()
loss = cls_loss
elif args.distill_model == 'simplified':
teacher_layer_num = len(teacher_atts)
student_layer_num = len(student_atts)
assert teacher_layer_num % student_layer_num == 0
layers_per_block = int(teacher_layer_num / student_layer_num)
new_teacher_atts = [
teacher_atts[i * layers_per_block + layers_per_block - 1]
for i in range(student_layer_num)
]
att_loss = 0.
rep_loss = 0.
# attention loss
for student_att, teacher_att in zip(student_atts,
new_teacher_atts):
student_att = torch.where(
student_att <= -1e2,
torch.zeros_like(student_att).to(device), student_att)
teacher_att = torch.where(
teacher_att <= -1e2,
torch.zeros_like(teacher_att).to(device), teacher_att)
tmp_loss = loss_mse(student_att, teacher_att)
att_loss += tmp_loss
# hidden states loss
new_teacher_reps = [
teacher_reps[i * layers_per_block]
for i in range(student_layer_num + 1)
]
new_student_reps = student_reps
for student_rep, teacher_rep in zip(new_student_reps,
new_teacher_reps):
tmp_loss = loss_mse(student_rep, teacher_rep)
rep_loss += tmp_loss
tr_att_loss += att_loss.item()
tr_rep_loss += rep_loss.item()
# classification loss
cls_loss = soft_loss + hard_loss
tr_cls_loss += cls_loss.item()
# total loss
loss = rep_loss + att_loss + cls_loss
else:
raise NotImplementedError
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
global_step += 1
if global_step % args.train_loss_step == 0:
loss = tr_loss / args.train_loss_step
cls_loss = tr_cls_loss / args.train_loss_step
att_loss = tr_att_loss / args.train_loss_step
rep_loss = tr_rep_loss / args.train_loss_step
loss_dict = {}
loss_dict['global_step'] = global_step
loss_dict['cls_loss'] = cls_loss
loss_dict['att_loss'] = att_loss
loss_dict['rep_loss'] = rep_loss
loss_dict['loss'] = loss
write_loss_to_file(loss_dict, output_loss_file)
tr_loss = 0.
tr_att_loss = 0.
tr_rep_loss = 0.
tr_cls_loss = 0.
if global_step % args.save_model_step == 0:
logger.info("***** Save model *****")
model_to_save = student_model.module if hasattr(
student_model, 'module') else student_model
model_name = WEIGHTS_NAME
checkpoint_name = 'checkpoint-' + str(global_step)
output_model_dir = os.path.join(args.output_dir,
checkpoint_name)
if not os.path.exists(output_model_dir):
os.makedirs(output_model_dir)
output_model_file = os.path.join(output_model_dir, model_name)
output_config_file = os.path.join(output_model_dir,
CONFIG_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(output_model_dir)
if os.path.exists(args.cache_file_dir):
import shutil
shutil.rmtree(args.cache_file_dir)