UIE迁移
UIE模型
UIE模型是指Unified Information Extraction,即通用信息抽取模型。该模型是由Yaojie Lu等人在ACL-2022中提出的通用信息抽取统一框架,实现了实体抽取、关系抽取、事件抽取、情感分析等任务的统一建模,并使得不同任务间具备良好的迁移和泛化能力。
模型迁移
一般模型迁移,需要经历以下步骤:
- 模型代码迁移
- 这一步需要进行
API
接口映射(比如把Pytorch
实现的模型代码映射成使用MindSpore
的API
接口实现,如果所需API
在目标框架中没有或者区别较大就可能需要自己实现了)
- 这一步需要进行
- 模型/模块验证
- 完成模型代码迁移后,必然少不了对其进行验证。在预训练权重迁移之前,如果模型结构比较简单,可以手动构造
Tensor
并且手动替代随机初始化权重等随机因素,然后将模型实例设置为预测模式,进行初步的输入输出验证(输出Tensor
的shape
应该相等、值应该误差在1e-5/1e-3/5e-3
内),这样可以验证出模型/模块的基本结构和前向计算过程是否正确。
- 完成模型代码迁移后,必然少不了对其进行验证。在预训练权重迁移之前,如果模型结构比较简单,可以手动构造
- 预训练权重迁移
- 当然最正确最有效的验证方式,是将模型的预训练权重加载到模型中,然后将模型设为预测模式,进行前向推理比较输出
Tensor
的shape
和值,并且一个完整的模型迁移过程通常本来就需要将预训练权重迁移也进行迁移。- 这时就出现一个问题,我们知道不同的深度学习框架除了
API
有所区别外,保存权重(也就是模型的参数)的格式也有所不同。该如何转换呢? - 实际上,模型的权重(参数)就是一堆数值矩阵,在模型结构一致的情况下可能出现的区别也就在权重矩阵的维度以及它们的名称上。因此,要实现迁移,只需要将它们的参数名字以及维度转化成适应目标框架的形式即可。(因此分析出两个框架之间的参数名称和权重维度对应关系非常重要!)
- 这时就出现一个问题,我们知道不同的深度学习框架除了
- 当然最正确最有效的验证方式,是将模型的预训练权重加载到模型中,然后将模型设为预测模式,进行前向推理比较输出
迁移过程
主要记录遇到的坑
框架
PaddleNLP->MindNLP
模型代码迁移
paddle.nn.MultiHeadAttention
与mindspore.nn.MultiheadAttention
paddle.nn.MultiHeadAttention
的attention 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] |
可以看到,对于
embeddings
和layer_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.weight
和pooler.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
88import 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