def test_compute(self):
head_importance, neuron_importance = compute_neuron_head_importance(
task_name='xnli',
model=self.model,
data_loader=self.data_loader,
num_layers=3,
num_heads=12)
assert (len(head_importance) == 3)
assert (len(neuron_importance) == 3)
self.reorder_neuron_head(self.model, head_importance, neuron_importance)
def do_train(args):
paddle.set_device("gpu" if args.n_gpu else "cpu")
if paddle.distributed.get_world_size() > 1:
paddle.distributed.init_parallel_env()
set_seed(args)
args.task_name = args.task_name.lower()
dataset_class, metric_class = TASK_CLASSES[args.task_name]
args.model_type = args.model_type.lower()
model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
train_ds = dataset_class.get_datasets(['train'])
tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
trans_func = partial(convert_example,
tokenizer=tokenizer,
label_list=train_ds.get_labels(),
max_seq_length=args.max_seq_length)
train_ds = train_ds.apply(trans_func, lazy=True)
train_batch_sampler = paddle.io.DistributedBatchSampler(
train_ds, batch_size=args.batch_size, shuffle=True)
batchify_fn = lambda samples, fn=Tuple(
Pad(axis=0, pad_val=tokenizer.pad_token_id), # input
Pad(axis=0, pad_val=tokenizer.pad_token_id), # segment
Stack(), # length
Stack(dtype="int64" if train_ds.get_labels() else "float32") # label
): [data for i, data in enumerate(fn(samples)) if i != 2]
train_data_loader = DataLoader(dataset=train_ds,
batch_sampler=train_batch_sampler,
collate_fn=batchify_fn,
num_workers=0,
return_list=True)
if args.task_name == "mnli":
dev_dataset_matched, dev_dataset_mismatched = dataset_class.get_datasets(
["dev_matched", "dev_mismatched"])
dev_dataset_matched = dev_dataset_matched.apply(trans_func, lazy=True)
dev_dataset_mismatched = dev_dataset_mismatched.apply(trans_func,
lazy=True)
dev_batch_sampler_matched = paddle.io.BatchSampler(
dev_dataset_matched, batch_size=args.batch_size, shuffle=False)
dev_data_loader_matched = DataLoader(
dataset=dev_dataset_matched,
batch_sampler=dev_batch_sampler_matched,
collate_fn=batchify_fn,
num_workers=0,
return_list=True)
dev_batch_sampler_mismatched = paddle.io.BatchSampler(
dev_dataset_mismatched, batch_size=args.batch_size, shuffle=False)
dev_data_loader_mismatched = DataLoader(
dataset=dev_dataset_mismatched,
batch_sampler=dev_batch_sampler_mismatched,
collate_fn=batchify_fn,
num_workers=0,
return_list=True)
else:
dev_dataset = dataset_class.get_datasets(["dev"])
dev_dataset = dev_dataset.apply(trans_func, lazy=True)
dev_batch_sampler = paddle.io.BatchSampler(dev_dataset,
batch_size=args.batch_size,
shuffle=False)
dev_data_loader = DataLoader(dataset=dev_dataset,
batch_sampler=dev_batch_sampler,
collate_fn=batchify_fn,
num_workers=0,
return_list=True)
num_labels = 1 if train_ds.get_labels() == None else len(
train_ds.get_labels())
model = model_class.from_pretrained(args.model_name_or_path,
num_classes=num_labels)
if paddle.distributed.get_world_size() > 1:
model = paddle.DataParallel(model)
# Step1: Initialize a dictionary to save the weights from the origin BERT model.
origin_weights = {}
for name, param in model.named_parameters():
origin_weights[name] = param
# Step2: Convert origin model to supernet.
sp_config = supernet(expand_ratio=args.width_mult_list)
model = Convert(sp_config).convert(model)
# Use weights saved in the dictionary to initialize supernet.
utils.set_state_dict(model, origin_weights)
del origin_weights
# Step3: Define teacher model.
teacher_model = model_class.from_pretrained(args.model_name_or_path,
num_classes=num_labels)
# Step4: Config about distillation.
mapping_layers = ['bert.embeddings']
for idx in range(model.bert.config['num_hidden_layers']):
mapping_layers.append('bert.encoder.layers.{}'.format(idx))
default_distill_config = {
'lambda_distill': 0.1,
'teacher_model': teacher_model,
'mapping_layers': mapping_layers,
}
distill_config = DistillConfig(**default_distill_config)
# Step5: Config in supernet training.
ofa_model = OFA(model,
distill_config=distill_config,
elastic_order=['width'])
criterion = paddle.nn.loss.CrossEntropyLoss() if train_ds.get_labels(
) else paddle.nn.loss.MSELoss()
metric = metric_class()
if args.task_name == "mnli":
dev_data_loader = (dev_data_loader_matched, dev_data_loader_mismatched)
# Step6: Calculate the importance of neurons and head,
# and then reorder them according to the importance.
head_importance, neuron_importance = utils.compute_neuron_head_importance(
args.task_name,
ofa_model.model,
dev_data_loader,
loss_fct=criterion,
num_layers=model.bert.config['num_hidden_layers'],
num_heads=model.bert.config['num_attention_heads'])
reorder_neuron_head(ofa_model.model, head_importance, neuron_importance)
lr_scheduler = paddle.optimizer.lr.LambdaDecay(
args.learning_rate,
lambda current_step, num_warmup_steps=args.warmup_steps,
num_training_steps=args.max_steps if args.max_steps > 0 else
(len(train_data_loader) * args.num_train_epochs): float(
current_step) / float(max(1, num_warmup_steps))
if current_step < num_warmup_steps else max(
0.0,
float(num_training_steps - current_step) / float(
max(1, num_training_steps - num_warmup_steps))))
optimizer = paddle.optimizer.AdamW(
learning_rate=lr_scheduler,
epsilon=args.adam_epsilon,
parameters=ofa_model.model.parameters(),
weight_decay=args.weight_decay,
apply_decay_param_fun=lambda x: x in [
p.name for n, p in ofa_model.model.named_parameters()
if not any(nd in n for nd in ["bias", "norm"])
])
global_step = 0
tic_train = time.time()
for epoch in range(args.num_train_epochs):
# Step7: Set current epoch and task.
ofa_model.set_epoch(epoch)
ofa_model.set_task('width')
for step, batch in enumerate(train_data_loader):
global_step += 1
input_ids, segment_ids, labels = batch
for width_mult in args.width_mult_list:
# Step8: Broadcast supernet config from width_mult,
# and use this config in supernet training.
net_config = apply_config(ofa_model, width_mult)
ofa_model.set_net_config(net_config)
logits, teacher_logits = ofa_model(input_ids,
segment_ids,
attention_mask=[None, None])
rep_loss = ofa_model.calc_distill_loss()
if args.task_name == 'sts-b':
logit_loss = 0.0
else:
logit_loss = soft_cross_entropy(logits,
teacher_logits.detach())
loss = rep_loss + args.lambda_logit * logit_loss
loss.backward()
optimizer.step()
lr_scheduler.step()
ofa_model.model.clear_gradients()
if global_step % args.logging_steps == 0:
if (not args.n_gpu > 1) or paddle.distributed.get_rank() == 0:
logger.info(
"global step %d, epoch: %d, batch: %d, loss: %f, speed: %.2f step/s"
% (global_step, epoch, step, loss, args.logging_steps /
(time.time() - tic_train)))
tic_train = time.time()
if global_step % args.save_steps == 0:
if args.task_name == "mnli":
evaluate(teacher_model,
criterion,
metric,
dev_data_loader_matched,
width_mult=100)
evaluate(teacher_model,
criterion,
metric,
dev_data_loader_mismatched,
width_mult=100)
else:
evaluate(teacher_model,
criterion,
metric,
dev_data_loader,
width_mult=100)
for idx, width_mult in enumerate(args.width_mult_list):
net_config = apply_config(ofa_model, width_mult)
ofa_model.set_net_config(net_config)
tic_eval = time.time()
if args.task_name == "mnli":
acc = evaluate(ofa_model, criterion, metric,
dev_data_loader_matched, width_mult)
evaluate(ofa_model, criterion, metric,
dev_data_loader_mismatched, width_mult)
print("eval done total : %s s" %
(time.time() - tic_eval))
else:
acc = evaluate(ofa_model, criterion, metric,
dev_data_loader, width_mult)
print("eval done total : %s s" %
(time.time() - tic_eval))
if (not args.n_gpu > 1
) or paddle.distributed.get_rank() == 0:
output_dir = os.path.join(args.output_dir,
"model_%d" % global_step)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# need better way to get inner model of DataParallel
model_to_save = model._layers if isinstance(
model, paddle.DataParallel) else model
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)