def reorder_neuron_head(self, model, head_importance, neuron_importance):
# reorder heads and ffn neurons
for layer, current_importance in enumerate(neuron_importance):
# reorder heads
idx = paddle.argsort(head_importance[layer], descending=True)
reorder_head(model.encoder.layers[layer].self_attn, idx)
# reorder neurons
idx = paddle.argsort(paddle.to_tensor(current_importance),
descending=True)
reorder_neuron(model.encoder.layers[layer].linear1, idx, dim=1)
reorder_neuron(model.encoder.layers[layer].linear2, idx, dim=0)
def reorder_neuron_head(model, head_importance, neuron_importance):
"""
Reorders weights according head importance and neuron importance
"""
from paddleslim.nas.ofa.utils import nlp_utils
# Reorders heads and ffn neurons
for layer, current_importance in enumerate(neuron_importance):
# Reorders heads
idx = paddle.argsort(head_importance[layer], descending=True)
nlp_utils.reorder_head(model.base_model.encoder.layers[layer].self_attn,
idx)
# Reorders neurons
idx = paddle.argsort(
paddle.to_tensor(current_importance), descending=True)
nlp_utils.reorder_neuron(
model.base_model.encoder.layers[layer].linear1.fn, idx, dim=1)
nlp_utils.reorder_neuron(
model.base_model.encoder.layers[layer].linear2.fn, idx, dim=0)