Example #1
0
 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)