def prune(args):
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Setup task, e.g., translation, language modeling, etc.
task = tasks.setup_task(args)
# Load dataset splits
load_dataset_splits(task, ['train', "valid"])
# Build model and criterion
model = task.build_model(args)
criterion = task.build_criterion(args)
print('| model {}, criterion {},'.format(args.arch,
criterion.__class__.__name__))
print('| num. model params: {}'.format(
sum(p.numel() for p in model.parameters())))
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
# Build trainer
trainer = Trainer(args, task, model, criterion, dummy_batch)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
print('| Optimizer {}'.format(trainer.optimizer.__class__.__name__))
# Initialize dataloader
epoch_itr = task.get_batch_iterator(
dataset=task.dataset(args.train_subset),
max_tokens=args.max_tokens,
max_sentences=args.max_sentences,
max_positions=max_positions,
ignore_invalid_inputs=True,
required_batch_size_multiple=8,
seed=args.seed,
num_shards=args.distributed_world_size,
shard_id=args.distributed_rank,
)
# Load the latest checkpoint if one is available
if not load_checkpoint(args, trainer, epoch_itr):
trainer.dummy_train_step([dummy_batch])
# Train until the learning rate gets too small
prune_meter = StopwatchMeter()
prune_meter.start()
# Estimate head importance scores
head_importance, head_stats = estimate_head_importance(
args, trainer, task, epoch_itr)
prune_meter.stop()
print('| done estimating head importance in {:.1f} seconds'.format(
prune_meter.sum))
torch.save(head_stats,
f"{os.path.dirname(args.restore_file)}/heads_stats.bin")
# Print
print("Head importances")
print("Encoder self attention")
for layer in range(head_importance["encoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_self"][layer]))
print("Encoder decoder attention")
for layer in range(head_importance["encoder_decoder"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_decoder"][layer]))
print("Decoder self attention")
for layer in range(head_importance["decoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["decoder_self"][layer]))
# Print sorted pruning profile
encoder_self_profile = get_profile(head_importance["encoder_self"],
prefix="E")
encoder_decoder_profile = get_profile(head_importance["encoder_decoder"],
prefix="A")
decoder_self_profile = get_profile(head_importance["decoder_self"],
prefix="D")
# Join all
all_profiles = {}
if not (args.decoder_self_only or args.encoder_decoder_only):
all_profiles.update(encoder_self_profile)
if not (args.encoder_self_only or args.decoder_self_only):
all_profiles.update(encoder_decoder_profile)
if not (args.encoder_self_only or args.encoder_decoder_only):
all_profiles.update(decoder_self_profile)
sorted_profiles = sorted(all_profiles.items(),
key=lambda x: x[1],
reverse=args.one_minus)
print("Heads sorted by importance:")
print(" ".join(p for p, _ in sorted_profiles))
print("Sorted head importance scores:")
print(" ".join(f"{v.data:.5f}" for _, v in sorted_profiles))
if args.only_importance:
return
tot_n_heads = len(sorted_profiles)
# Eval pruning
if args.one_head:
kept_layers = set()
to_prune_profile = []
for p, _ in reversed(sorted_profiles):
layer_name = ":".join(p.split(":")[:-1])
if layer_name not in kept_layers:
kept_layers.add(layer_name)
continue
else:
to_prune_profile.insert(0, p)
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
return
for i in range(0, 10):
n_to_prune = int(ceil(tot_n_heads * i / 10))
to_prune_profile = [p for p, _ in sorted_profiles[:n_to_prune]]
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()
def prune2(args, task, model, trainer, epoch_itr): # changed2
if args.max_tokens is None:
args.max_tokens = 6000
print(args)
# avoid aliasing
task = copy.deepcopy(task)
model = copy.deepcopy(model)
trainer = copy.deepcopy(trainer)
epoch_itr = copy.deepcopy(trainer)
if not torch.cuda.is_available():
raise NotImplementedError('Training on CPU is not supported')
torch.cuda.set_device(args.device_id)
torch.manual_seed(args.seed)
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
# placeholder DistributedDataParallel when there's an uneven number of
# batches per worker.
max_positions = utils.resolve_max_positions(
task.max_positions(),
model.max_positions(),
)
dummy_batch = task.dataset('train').get_dummy_batch(
args.max_tokens, max_positions)
print('| training on {} GPUs'.format(args.distributed_world_size))
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
args.max_tokens,
args.max_sentences,
))
print('| Optimizer {}'.format(trainer.optimizer.__class__.__name__))
# Train until the learning rate gets too small
prune_meter = StopwatchMeter()
prune_meter.start()
# Estimate head importance scores
head_importance, head_stats = estimate_head_importance(
args, trainer, task, epoch_itr)
prune_meter.stop()
print('| done estimating head importance in {:.1f} seconds'.format(
prune_meter.sum))
torch.save(head_stats,
f"{os.path.dirname(args.restore_file)}/heads_stats.bin")
# Print
print("Head importances")
print("Encoder self attention")
for layer in range(head_importance["encoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_self"][layer]))
print("Encoder decoder attention")
for layer in range(head_importance["encoder_decoder"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["encoder_decoder"][layer]))
print("Decoder self attention")
for layer in range(head_importance["decoder_self"].size(0)):
print("\t".join(f"{x:.5f}"
for x in head_importance["decoder_self"][layer]))
# Print sorted pruning profile
encoder_self_profile = get_profile(head_importance["encoder_self"],
prefix="E")
encoder_decoder_profile = get_profile(head_importance["encoder_decoder"],
prefix="A")
decoder_self_profile = get_profile(head_importance["decoder_self"],
prefix="D")
# Join all
all_profiles = {}
if not (args.decoder_self_only or args.encoder_decoder_only):
all_profiles.update(encoder_self_profile)
if not (args.encoder_self_only or args.decoder_self_only):
all_profiles.update(encoder_decoder_profile)
if not (args.encoder_self_only or args.encoder_decoder_only):
all_profiles.update(decoder_self_profile)
sorted_profiles = sorted(all_profiles.items(),
key=lambda x: x[1],
reverse=args.one_minus)
print("Heads sorted by importance:")
print(" ".join(p for p, _ in sorted_profiles))
print("Sorted head importance scores:")
print(" ".join(f"{v.data:.5f}" for _, v in sorted_profiles))
tot_n_heads = len(sorted_profiles)
for i in range(0, 10):
n_to_prune = int(ceil(tot_n_heads * i / 10))
to_prune_profile = [p for p, _ in sorted_profiles[:n_to_prune]]
to_prune = parse_head_pruning_descriptors(to_prune_profile,
reverse_descriptors=False)
print(f"Evaluating following profile: \t{' '.join(to_prune_profile)}")
# Apply pruning
mask_heads(model, to_prune, args.transformer_mask_rescale)
bleu = eval_bleu_score(
model,
task,
task.dataset(args.valid_subset),
beam=args.beam,
replace_unk=args.replace_unk,
lenpen=args.lenpen,
buffer_size=100,
use_cuda=torch.cuda.is_available() and not args.cpu,
remove_bpe=args.remove_bpe,
max_sentences=args.max_sentences,
max_tokens=args.max_tokens,
stop_early=not args.no_early_stop,
normalize_scores=not args.unnormalized,
min_len=args.min_len,
)
print(f"BLEU score: \t{bleu.score:.2f}")
sys.stdout.flush()