UIE迁移

UIE迁移

UIE模型

UIE模型是指Unified Information Extraction,即通用信息抽取模型。该模型是由Yaojie Lu等人在ACL-2022中提出的通用信息抽取统一框架,实现了实体抽取、关系抽取、事件抽取、情感分析等任务的统一建模,并使得不同任务间具备良好的迁移和泛化能力。

模型迁移

一般模型迁移,需要经历以下步骤:

  • 模型代码迁移
    • 这一步需要进行API接口映射(比如把Pytorch实现的模型代码映射成使用MindSporeAPI接口实现,如果所需API在目标框架中没有或者区别较大就可能需要自己实现了)
  • 模型/模块验证
    • 完成模型代码迁移后,必然少不了对其进行验证。在预训练权重迁移之前,如果模型结构比较简单,可以手动构造Tensor并且手动替代随机初始化权重等随机因素,然后将模型实例设置为预测模式,进行初步的输入输出验证(输出Tensorshape应该相等、值应该误差在1e-5/1e-3/5e-3内),这样可以验证出模型/模块的基本结构和前向计算过程是否正确。
  • 预训练权重迁移
    • 当然最正确最有效的验证方式,是将模型的预训练权重加载到模型中,然后将模型设为预测模式,进行前向推理比较输出Tensorshape和值,并且一个完整的模型迁移过程通常本来就需要将预训练权重迁移也进行迁移。
      • 这时就出现一个问题,我们知道不同的深度学习框架除了API有所区别外,保存权重(也就是模型的参数)的格式也有所不同。该如何转换呢?
      • 实际上,模型的权重(参数)就是一堆数值矩阵,在模型结构一致的情况下可能出现的区别也就在权重矩阵的维度以及它们的名称上。因此,要实现迁移,只需要将它们的参数名字以及维度转化成适应目标框架的形式即可。(因此分析出两个框架之间的参数名称和权重维度对应关系非常重要!)

迁移过程

主要记录遇到的

框架

  • PaddleNLP->MindNLP

模型代码迁移

  • paddle.nn.MultiHeadAttentionmindspore.nn.MultiheadAttention

    • paddle.nn.MultiHeadAttentionattention mask的输入维度:[batch_size,n_head,sequence_length,sequence_length]

    • mindspore.nn.MultiheadAttention的的attention mask的输入维度: [batch_size * n_head, sequence_length,sequence_length]

      尾部两个sequence_length维度分别代表源序列长度和目标序列长度

    • 因此在MindNLP中需要对attention mask加一步tile操作, 将Tensor维度变成第二种形式。

      1
      2
      3
      4
      5
      6
      7
      8
      9
      # paddle 中的 attention mask
      attention_mask = paddle.unsqueeze(
      (input_ids == self.pad_token_id).astype(self.pooler.dense.weight.dtype) * -1e4, axis=[1, 2]
      )
      # mindspore 中的 attention mask
      attention_mask = ((input_ids == self.pad_token_id).astype(self.pooler.dense.weight.dtype) * -1e4).unsqueeze(1).unsqueeze(2)
      attention_mask = ops.tile(attention_mask, (1, self.nheads, seq_length, 1)).reshape(-1, seq_length, seq_length)

      # 注意:mindspore 的 unsqueeze 一次只支持增加一维,但可以使用链式调用

权重迁移

  • 坑同样出现在MultiHeadAttention当中,首先来看看源码:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    # paddle MultiHeadAttention 中参数定义

    self.q_proj = Linear(embed_dim, embed_dim, weight_attr, bias_attr=bias_attr)
    self.k_proj = Linear(self.kdim, embed_dim, weight_attr, bias_attr=bias_attr)
    self.v_proj = Linear(self.vdim, embed_dim, weight_attr, bias_attr=bias_attr)

    # mindspore MultiHeadAttention 中参数定义
    if not self._qkv_same_embed_dim:
    self.q_proj_weight = Parameter(initializer(XavierUniform(), (embed_dim, embed_dim)), 'q_proj_weight')
    self.k_proj_weight = Parameter(initializer(XavierUniform(), (embed_dim, self.kdim)), 'k_proj_weight')
    self.v_proj_weight = Parameter(initializer(XavierUniform(), (embed_dim, self.vdim)), 'v_proj_weight')
    self.in_proj_weight = None
    else:
    self.in_proj_weight = Parameter(initializer(XavierUniform(), (3 * embed_dim, embed_dim)), 'in_proj_weight')
    self.q_proj_weight = None
    self.k_proj_weight = None
    self.v_proj_weight = None
    • 首先可以看到,在paddle中,qkv参数是直接用全连接层定义的,而在mindspore中使用的是Parameter,当然实际上两种方式的效果一样所以这一点区别不重要。
    • 重点在于,在mindspore中,如果qkv的维度相同的话,它会把原本的qkv对应的三个参数矩阵拼接成一个in_proj_weight,也就是esle部分的代码所做的事。因此,在checkpoint转化时,也需要将这三个参数矩阵进行拼接融合成一个。

    • 下面分别打印在两个框架中实现的模型的参数名及权重维度:

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      14
      15
      16
      17
      18
      19
      20
      21
      22
      23
      24
      25
      26
      27
      28
      29
      30
      31
      32
      33
      34
      35
      36
      37
      38
      39
      40
      41
      42
      43
      44
      45
      46
      47
      48
      49
      50
      51
      52
      53
      54
      55
      56
      57
      # paddle实现UIE的权重
      ernie.embeddings.word_embeddings.weight[40000, 768]
      ernie.embeddings.position_embeddings.weight[2048, 768]
      ernie.embeddings.token_type_embeddings.weight[4, 768]
      ernie.embeddings.task_type_embeddings.weight[3, 768]
      ernie.embeddings.layer_norm.weight[768]
      ernie.embeddings.layer_norm.bias[768]
      ernie.encoder.layers.0.self_attn.q_proj.weight[768, 768]
      ernie.encoder.layers.0.self_attn.q_proj.bias[768]
      ernie.encoder.layers.0.self_attn.k_proj.weight[768, 768]
      ernie.encoder.layers.0.self_attn.k_proj.bias[768]
      ernie.encoder.layers.0.self_attn.v_proj.weight[768, 768]
      ernie.encoder.layers.0.self_attn.v_proj.bias[768]
      ernie.encoder.layers.0.self_attn.out_proj.weight[768, 768]
      ernie.encoder.layers.0.self_attn.out_proj.bias[768]
      ernie.encoder.layers.0.linear1.weight[768, 3072]
      ernie.encoder.layers.0.linear1.bias[3072]
      ernie.encoder.layers.0.linear2.weight[3072, 768]
      ernie.encoder.layers.0.linear2.bias[768]
      ernie.encoder.layers.0.norm1.weight[768]
      ernie.encoder.layers.0.norm1.bias[768]
      ernie.encoder.layers.0.norm2.weight[768]
      ernie.encoder.layers.0.norm2.bias[768]
      ...
      ernie.pooler.dense.weight[768, 768]
      ernie.pooler.dense.bias[768]
      linear_start.weight[768, 1]
      linear_start.bias[1]
      linear_end.weight[768, 1]
      linear_end.bias[1]

      # mindspore实现UIE的权重
      ernie.embeddings.word_embeddings.embedding_table[40000, 768]
      ernie.embeddings.position_embeddings.embedding_table[2048, 768]
      ernie.embeddings.token_type_embeddings.embedding_table[4, 768]
      ernie.embeddings.task_type_embeddings.embedding_table[3, 768]
      ernie.embeddings.layer_norm.gamma[768]
      ernie.embeddings.layer_norm.beta[768]
      ernie.encoder.layers.0.self_attn.in_proj_weight[2304, 768]
      ernie.encoder.layers.0.self_attn.in_proj_bias[2304]
      ernie.encoder.layers.0.self_attn.out_proj.weight[768, 768]
      ernie.encoder.layers.0.self_attn.out_proj.bias[768]
      ernie.encoder.layers.0.linear1.weight[3072, 768]
      ernie.encoder.layers.0.linear1.bias[3072]
      ernie.encoder.layers.0.linear2.weight[768, 3072]
      ernie.encoder.layers.0.linear2.bias[768]
      ernie.encoder.layers.0.norm1.gamma[768]
      ernie.encoder.layers.0.norm1.beta[768]
      ernie.encoder.layers.0.norm2.gamma[768]
      ernie.encoder.layers.0.norm2.beta[768]
      ...
      ernie.pooler.dense.weight[768, 768]
      ernie.pooler.dense.bias[768]
      linear_start.weight[1, 768]
      linear_start.bias[1]
      linear_end.weight[1, 768]
      linear_end.bias[1]
  • 为了方便分析,...省略了11层ernie.encoder.layers(总共12层)。下面表格列出两个框架之间的参数区别(不同之处已加粗):

paddle mindspore
word_embeddings.weight[40000, 768] word_embeddings.embedding_table[40000, 768]
position_embeddings.weigh[2048, 768] word_embeddings.embedding_table[2048, 768]
token_type_embeddings.weight[4, 768] token_type_embeddings.embedding_table[4, 768]
layer_norm.weight[768] layer_norm.gamma[768]
layer_norm.bias[768] layer_norm.beta[768]
layers.0.self_attn.q_proj.weight[768, 768]
layers.0.self_attn.q_proj.bias[768]
layers.0.self_attn.k_proj.weight[768, 768]
layers.0.self_attn.k_proj.bias[768]
layers.0.self_attn.v_proj.weight[768, 768]
layers.0.self_attn.v_proj.bias[768]
layers.0.self_attn.in_proj_weight[2304, 768]
layers.0.self_attn.in_proj_bias[2304]
layers.0.linear1.weight[768, 3072] layers.0.linear1.weight[3072, 768]
layers.0.linear2.weight[3072, 768] layers.0.linear2.weight[768, 3072]
layers.0.norm1.weight[768] layers.0.norm1.gamma[768]
layers.0.norm1.bias[768] layers.0.norm1.beta[768]
layers.0.norm2.weight[768] layers.0.norm2.gamma[768]
layers.0.norm2.bias[768] layers.0.norm2.beta[768]
linear_start.weight[768, 1] linear_start.weight[1, 768]
linear_end.weight[768, 1] linear_end.weight[1, 768]
  • 可以看到,对于embeddingslayer_norm/norm区别在于参数名而维度形状相同;对于全连接层liner区别在于维度形状交换而参数名相同;再看表格的第6行,前面源代码已经分析过了,paddle使用了全连接层并且没有将qkv的参数融合在一起,所以它的qkv的权重矩阵刚好是mindspore中权重矩阵的拆分(这里需要注意拼接weight的时候先沿axis=1拼接成[760,2304],然后再转置/交换维度成[2304,768],如果沿axis=0直接拼接,qkv所对应的权重矩阵就乱了)。

  • 好到这里就分析完了?未列入上表的其他参数说明一模一样,即参数名相同维度也没反?

  • 注意:大坑来了!!!

paddle mindspore
layers.0.self_attn.out_proj.weight [768, 768] layers.0.self_attn.out_proj.weight [768, 768]
pooler.dense.weight [768, 768] pooler.dense.weight [768, 768]
  • 特意把这self_attn.out_proj.weightpooler.dense.weight单独提出来,这两个参数乍一看在两个框架中完全一样!无论是参数名称还是其维度!但是,实际上是不同的!两个矩阵虽然在维度大小上完全一致,但是有一个转置关系!也就说这两个参数矩阵也需要交换一下维度!

  • 经过上面分析,接下来要做的就是:1. 将参数名进行转换 2. 将维度进行转换

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    import logging
    import re
    import numpy as np
    import paddle
    from mindspore import Tensor
    from mindspore.train.serialization import save_checkpoint

    def paddle_to_mindspore(pth_file, size:str=None):

    size = "mindspore" if not size else size # rename ckpt

    logging.info('Starting checkpoint conversion.')
    ms_ckpt = []
    state_dict = paddle.load("model_state.pdparams")

    cnt1 = 0
    in_p_w = np.empty((768,0),dtype=np.float32) # 创建一个维度与embedding相同的空矩阵以便拼接
    cnt2 = 0
    in_p_b = np.empty(0,dtype=np.float32)

    for k, v in state_dict.items():
    if 'embeddings.weight' in k:
    k = k.replace('embeddings.weight', 'embeddings.embedding_table')
    ms_ckpt.append({'name': k, 'data': Tensor(v.numpy())})
    continue

    pattern = r"norm.*?\.weight"
    if re.search(pattern, k):
    k = re.sub("weight", "gamma", k)
    ms_ckpt.append({'name': k, 'data': Tensor(v.numpy())})
    continue

    pattern = r"norm.*?\.bias"
    if re.search(pattern, k):
    k = re.sub("bias", "beta", k)
    ms_ckpt.append({'name': k, 'data': Tensor(v.numpy())})
    continue

    pattern = r"linear.*\.weight"
    if re.search(pattern, k):
    ms_ckpt.append({'name': k, 'data': Tensor(np.transpose(v.numpy()))})
    continue

    pattern = r"self_attn\.._proj.weight"
    if re.search(pattern, k):
    k = re.sub(r"self_attn\.._proj.weight", "self_attn.in_proj_weight", k)
    if cnt1 % 3 != 0 or cnt1 == 0:
    in_p_w = np.concatenate((in_p_w, v.numpy()), axis=1)
    if (cnt1+1) % 3 == 0:
    ms_ckpt.append({'name': k, 'data': Tensor(in_p_w).swapaxes(0,1)})
    else: # 每三个拼接一次(qkv)对齐 mindspore 的 in_proj_weight
    in_p_w = v.numpy()
    cnt1 += 1
    continue

    pattern = r"self_attn\.._proj.bias"
    if re.search(pattern, k):
    k = re.sub(r"self_attn\.._proj.bias", "self_attn.in_proj_bias", k)
    if cnt2 % 3 != 0 or cnt2 == 0:
    in_p_b = np.append(in_p_b, v.numpy())
    if (cnt2+1) % 3 == 0:
    ms_ckpt.append({'name': k, 'data': Tensor(in_p_b)})
    else:
    in_p_b = v.numpy()
    cnt2 += 1
    continue

    # 很关键,padlle权重维度与mindspore的权重维度很多都是反的,特别注意多个维度相等的时候也要转换!
    pattern = r"self_attn\.out_proj.weight"
    if re.search(pattern, k):
    ms_ckpt.append({'name': k, 'data': Tensor(np.transpose(v.numpy()))})
    continue

    pattern = r"pooler\.dense.weight"
    if re.search(pattern, k):
    ms_ckpt.append({'name': k, 'data': Tensor(np.transpose(v.numpy()))})
    continue

    ms_ckpt.append({'name': k, 'data': Tensor(v.numpy())})


    ms_ckpt_path = pth_file.replace('.pdparams','.ckpt')
    try:
    save_checkpoint(ms_ckpt, ms_ckpt_path)
    except:
    raise RuntimeError(f'Save checkpoint to {ms_ckpt_path} failed, please checkout the path.')

    return ms_ckpt_path
作者

TIANYUZHOU

发布于

2023-04-05

更新于

2023-04-18

许可协议

评论