关于Transformer中position在代码实现的一点思考

Transformer需要显式加position,不同版本的代码有不同创建position的方式。前几天在实现中遇到了一些小问题,因此记录在此。

过去在获取position的时候,一般都是在外部创建,也即在做batch的时候创建。具体而言:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def collate_fn(batch):
sentences, labels = [], []
for b in batch:
sentences.append(b[0])
labels.append(b[1])

max_len = max(len(sent) for sent in sentences)
batch_seq = np.array([
sent + [0] * (max_len - len(sent))
for sent in sentences
])

batch_pos = np.array([
[pos_i + 1 if word_i != 0 else 0
for pos_i, word_i in enumerate(sent)] for sent in batch_seq
])

batch_label = np.array(labels)

batch_seq = torch.LongTensor(batch_seq)
batch_pos = torch.LongTensor(batch_pos)
batch_label = torch.LongTensor(batch_label)

return batch_seq, batch_pos, batch_label

collate_fn是Dataloader的参数,具体的作用是在获取了sample和label后显式padding以及转换为tensor等操作。

而在模型训练的时候,则将position作为参数显式传入forward里:

1
model(batch_seq, batch_pos)

这是一种方法,但我认为看起来不大优雅。假设以下情形:我希望将Transformer作为encoder来分类,同时希望与LSTM作为encoder来对比。但LSTM不需要position,且position作为参数已经写死在函数里了,这样一来,就不能很方便地换模块了。

因此,我希望能够在模型内部创建position,使得传入的参数只有batch_seq,这样更加优雅。

我一开始的办法是封装一个EmbeddingLayer

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
class EmbeddingLayer(nn.Module):
def __init__(self, args, pretrained_matrix=None):
super(EmbeddingLayer, self).__init__()
self.word_embedding = nn.Embedding(args.vocab_size, args.embed_dim, padding_idx=0)
self.embed_factor = args.embed_dim ** (1 / 2)
self.dropout = nn.Dropout(args.dropout)
if pretrained_matrix is not None:
pretrained_matrix = torch.from_numpy(pretrained_matrix).type(torch.FloatTensor)
self.word_embedding.weight = nn.Parameter(pretrained_matrix, requires_grad=True)

n_position = args.max_sent_len + 1
self.position_embedding = nn.Embedding.from_pretrained(
get_sinusoid_encoding_table(n_position, args.embed_dim, padding_idx=0),
freeze=True)

def forward(self, seq):
x = self.word_embedding(seq) * self.embed_factor

pos = self._create_pos(seq)
x = x + self.position_embedding(pos)

return self.dropout(x)

def _create_pos(self, seq):
"""
:param seq: batch,seq
:return:
"""
seq_pos = np.array([
[pos_i + 1 if word_i != 0 else 0
for pos_i, word_i in enumerate(sent)] for sent in seq
])
return seq_pos.to(seq.device)

也即在EmbeddingLayer内部根据seq来生成position。

虽然逻辑上是正确的,但看GPU的利用率总是很低。排查了很久之后,才发现上述_create_pos(self, seq)的效率太低了。

看起来似乎没什么问题,这和collate_fn的做法是一样的。但为什么就是效率那么低?

思考了一下,我推测是因为,model内部一般都是tensor操作,一般都是在GPU上的操作;而_create_pos作为model内部的函数,却是做了在CPU上的操作。这样一来,本来整个GPU流水线的运算,却要等这个CPU的操作完成了才能往下做。这就导致了利用率低下的问题。

为了验证我的想法,我尝试将代码改成tensor的操作:

1
2
3
4
5
6
7
8
def _create_pos(self, seq):
"""

:param seq: batch,seq
:return:
"""
batch_size, max_seq_len = seq.size()
pos_tensor = torch.arange(1, max_seq_len + 1)[None, :].expand(batch_size, -1).long() # batch,max_seq_len we don't do padding

首先上述代码是完全的tensor操作,是在GPU上完成的;第二,我意识到其实不需要显式将sequence中padding的位置对应的position也置为0。因为只要padding的位置不参与计算loss就没有关系;或者将padding的操作交给后续的mask来做即可。

这样改一下,发现GPU利用率终于大幅提升,基本上95%以上。

那么另一问题又来了,为啥之前collate_fn的做法,也即在外部创建pos没问题,同样的做法移到model内部来就有问题?

我推测,第一,在CPU上的操作是可以多线程的,这边CPU边创建新的batch,另一边model在GPU边训练,这两个(CPU和GPU)实质上是并行的而不是串行的;第二,如果将其移到内部操作,就变成了外部CPU创建好batch,传入model后,model内部还要经历 GPU —> CPU —> GPU的操作,显然CPU会拖累整个进度,且导致利用率低下的问题。


就这么一个小小的问题,我就debug了一天,一开始怀疑是我自己的iterator写的效率低,然后换成了官方的Dataloader且开了10个进程,但仍然有问题;后来怀疑是模型写错了,check了半天都没有发现哪里有问题;然后是开Pycharm的profile查看是哪些函数占用太多时间,但因为_create_pos本身是没有运行很长时间的,只是阻碍了GPU的运行效率,因此也没有发现这个问题;也是到最后,尝试不用position后,才发现效率有大幅提升。

从这次的debug经历,有两个深刻教训:①model内部尽可能利用GPU而不要做太多的CPU操作,否则就可能出现上述问题;②在debug要大胆尝试控制变量,而不要仅仅是肉眼去看代码找bug。