torch.gather()使用方法
问题分析
在阅读OneIE代码时，突然看到一段代码十分精妙，用来预测BERT等预训练语言模型在使用tokenizer进行分词时，会将一个单词可能分成多个token，如原始句子为"(END VIDEO CLIP)"，正常理解按照空格和字符进行划分为["(", "END", "VIDEO", "CLIP", ")"]，然而BERT的划分结果为["(", "E", "##ND", "VI", "##DE", "##O", "C", "##L", "##IP", ")"]，则每个token对应的长度为token_lens=[1, 2, 3, 3, 1]。所以在经过BERT之后，单词END的表征将变成了两个部分(E,##ND)，所以END的表示便有了歧义，常见的做法有：1、取E和##ND表征的平均值。2、取E来表示END。3、E和##ND表征进行连接，在经过线性变化进行降维。如果还有其他做法可以进行补充。
为什么上述问题要用到torch.gather()函数呢？这个是因为这个函数可以按照给出的索引对原始的tensor进行取值，类似于列表中的索引和切片。
原理分析
先看官方文档，官方文档给出了函数的定义及其相关解析。
红色框中表明了函数根据index的索引进行取值的规则。建议多看几遍！！！dim的取值为多少，就代表
函数存在三个输入参数：
input:表示输入向量 dim:按照该轴进行取值，和常规的函数相同用法 index:需要在输入向量中取值的索引位置
值得注意是，dim的值要小于输入向量input的维度，如果是一个二维的向量，则dim只能取值为0或1，和常规的函数相同使用。函数输出的向量形状和index向量必须一致。index向量中的取值要小于input的形状维度。更加详细的规则如下所示：
以官方示例为例，
扩展dim=0将会产生不一样的结果
BERT的token表征分析
回到问题中，部分token在tokenizer时会分成多个表示，我们将以OneIE代码为例进行分析和讲解：
from transformers import BertTokenizer, BertModel import torch def token_lens_to_idxs(token_lens): """Map token lengths to a word piece index matrix (for torch.gather) and a mask tensor. For example (only show a sequence instead of a batch): token lengths: [1,1,1,3,1] => indices: [[0,0,0], [1,0,0], [2,0,0], [3,4,5], [6,0,0]] masks: [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.33, 0.33, 0.33], [1.0, 0.0, 0.0]] Next, we use torch.gather() to select vectors of word pieces for each token, and average them as follows (incomplete code): outputs = torch.gather(bert_outputs, 1, indices) * masks outputs = bert_outputs.view(batch_size, seq_len, -1, self.bert_dim) outputs = bert_outputs.sum(2) :param token_lens (list): token lengths. (batch,seq_len) :return: a index matrix and a mask tensor. """ max_token_num = max([len(x) for x in token_lens]) max_token_len = max([max(x) for x in token_lens]) idxs, masks = [], [] for seq_token_lens in token_lens: seq_idxs, seq_masks = [], [] offset = 0 for token_len in seq_token_lens: seq_idxs.extend([i + offset for i in range(token_len) ] + [-1] * (max_token_len - token_len)) seq_masks.extend([1.0 / token_len] * token_len + [0.0] * (max_token_len - token_len)) offset += token_len seq_idxs.extend([-1] * max_token_len * (max_token_num - len(seq_token_lens))) seq_masks.extend([0.0] * max_token_len * (max_token_num - len(seq_token_lens))) idxs.append(seq_idxs) masks.append(seq_masks) return idxs, masks, max_token_num, max_token_len def data_process(datas): token_lens = [] tokens = [] pieces = [] sentences = [] for data in datas: token_lens.append(data['token_lens']) tokens.append(data['tokens']) pieces.append(data['pieces']) sentences.append(data['sentence']) return token_lens, tokens, pieces, sentences def get_bert_input(pieces, tokenizer, max_length=24): _piece_idxs = [] _attn_masks = [] for piece in pieces: piece_idxs = tokenizer.encode(piece, add_special_tokens=True, max_length=max_length, truncation=True) pad_num = max_length - len(piece_idxs) attn_mask = [1] * len(piece_idxs) + [0] * pad_num piece_idxs = piece_idxs + [0] * pad_num _piece_idxs.append(piece_idxs) _attn_masks.append(attn_mask) _piece_idxs = torch.LongTensor(_piece_idxs) _attn_masks = torch.LongTensor(_attn_masks) return _piece_idxs, _attn_masks data_example = [{"doc_id": "CNN_IP_20030409.1600.02", "sent_id": "CNN_IP_20030409.1600.02-21", "tokens": ["Yet", "until", "this", "war", "is", "fully", "won", ",", "we", "cannot", "be", "overconfident", "in", "our", "position", "."], "pieces": ["Yet", "until", "this", "war", "is", "fully", "won", ",", "we", "cannot", "be", "over", "##con", "##fi", "##dent", "in", "our", "position", "."], "token_lens": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 1], "sentence": "Yet until this war is fully won, we cannot be overconfident in our position.", "entity_mentions": [{"id": "CNN_IP_20030409.1600.02-E10-53", "text": "we", "entity_type": "GPE", "mention_type": "PRO", "entity_subtype": "Nation", "start": 8, "end": 9}, {"id": "CNN_IP_20030409.1600.02-E10-54", "text": "our", "entity_type": "GPE", "mention_type": "PRO", "entity_subtype": "Nation", "start": 13, "end": 14}], "relation_mentions": [], "event_mentions": [{"id": "CNN_IP_20030409.1600.02-EV1-1", "event_type": "Conflict:Attack", "trigger": {"text": "war", "start": 3, "end": 4}, "arguments": []}]}, {"doc_id": "CNN_IP_20030409.1600.02", "sent_id": "CNN_IP_20030409.1600.02-22", "tokens": ["And", "we", "must", "not", "underestimate", "the", "desperation", "of", "whatever", "forces", "remain", "loyal", "to", "the", "dictator", "."], "pieces": ["And", "we", "must", "not", "under", "##est", "##imate", "the", "desperation", "of", "whatever", "forces", "remain", "loyal", "to", "the", "dictator", "."], "token_lens": [1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "sentence": "And we must not underestimate the desperation of whatever forces remain loyal to the dictator.", "entity_mentions": [{"id": "CNN_IP_20030409.1600.02-E10-55", "text": "we", "entity_type": "GPE", "mention_type": "PRO", "entity_subtype": "Nation", "start": 1, "end": 2}, { "id": "CNN_IP_20030409.1600.02-E31-56", "text": "forces", "entity_type": "PER", "mention_type": "NOM", "entity_subtype": "Group", "start": 9, "end": 10}, {"id": "CNN_IP_20030409.1600.02-E12-57", "text": "dictator", "entity_type": "PER", "mention_type": "NOM", "entity_subtype": "Individual", "start": 14, "end": 15}], "relation_mentions": [{"id": "CNN_IP_20030409.1600.02-R17-1", "relation_type": "GEN-AFF", "relation_subtype": "GEN-AFF:Citizen-Resident-Religion-Ethnicity", "arguments": [{"entity_id": "CNN_IP_20030409.1600.02-E31-56", "text": "forces", "role": "Arg-1"}, {"entity_id": "CNN_IP_20030409.1600.02-E12-57", "text": "dictator", "role": "Arg-2"}]}], "event_mentions": []}, {"doc_id": "CNN_IP_20030409.1600.02", "sent_id": "CNN_IP_20030409.1600.02-23", "tokens": ["(", "END", "VIDEO", "CLIP", ")"], "pieces": ["(", "E", "##ND", "VI", "##DE", "##O", "C", "##L", "##IP", ")"], "token_lens": [ 1, 2, 3, 3, 1], "sentence": "(END VIDEO CLIP)", "entity_mentions": [], "relation_mentions": [], "event_mentions": []}, {"doc_id": "CNN_IP_20030409.1600.02", "sent_id": "CNN_IP_20030409.1600.02-18", "tokens": ["(", "BEGIN", "VIDEO", "CLIP", ")"], "pieces": ["(", "B", "##EG", "##IN", "VI", "##DE", "##O", "C", "##L", "##IP", ")"], "token_lens": [1, 3, 3, 3, 1], "sentence": "(BEGIN VIDEO CLIP)", "entity_mentions": [], "relation_mentions": [], "event_mentions": []}] bert_dim = 1536 batch_size = 1 tokenizer = BertTokenizer.from_pretrained('./bert/bert-base-cased') # output_hidden_states=True，表示输出bert中间层的结果 bert = BertModel.from_pretrained( './bert/bert-base-cased', output_hidden_states=True) token_lens, tokens, pieces, sentences = data_process(data_example) piece_idxs, attention_masks = get_bert_input(pieces, tokenizer, 24) all_bert_outputs = bert(piece_idxs, attention_mask=attention_masks) bert_outputs = all_bert_outputs[0] # 取BERT倒数第三层的输出连接，使效果更佳 extra_bert_outputs = all_bert_outputs[2][-3] bert_outputs = torch.cat([bert_outputs, extra_bert_outputs], dim=2) # 最为关键多个token融合，并选出最终的bert输出表示 idxs, masks, token_num, token_len = token_lens_to_idxs(token_lens) # +1 是因为第一个向量是[CLS]，并且将idxs中最小值有-1变化为0，expand是为了进行广播 idxs = piece_idxs.new(idxs).unsqueeze(-1).expand(batch_size, -1, bert_dim) + 1 # 便于后续的矩阵逐元素乘法 masks = bert_outputs.new(masks).unsqueeze(-1) # 逐元素乘法，因为mask会对多个词token进行平均化，并且将没有分割的token的mask填充置0 bert_outputs = torch.gather(bert_outputs, 1, idxs) * masks bert_outputs = bert_outputs.view(batch_size, token_num, token_len, bert_dim) bert_outputs = bert_outputs.sum(2) print(bert_outputs)
上述代码就是采用方法一，取E和##ND表征的平均值进行特征融合，如果还是很复杂，可以看以下精简版代码示例就可以看出特征融合的原理，最好自己运行代码查看：
import torch bert_outputs = torch.rand((1, 12, 8)) idxs = [0, -1, -1, 1, -1, -1, 2, 3, -1, 4, -1, -1] masks = [1, 0, 0, 1, 0, 0, 0.5, 0.5, 0, 1, 0, 0] idxs = torch.LongTensor(idxs) idxs = idxs.unsqueeze(-1).expand(1, -1, 8) + 1 # 便于后续的矩阵逐元素乘法 masks = bert_outputs.new(masks).unsqueeze(-1) # 逐元素乘法，因为mask会对多个词token进行平均化，并且将没有分割的token的mask填充置0, # dim=1表示在bert_outputs的seq上进行采样，gather对batch和dimension上维度保持不变，在乘mask则会将多余的清零 print(bert_outputs) bert_outputs = torch.gather(bert_outputs, 1, idxs) print(bert_outputs) bert_outputs=bert_outputs* masks print(bert_outputs) bert_outputs = bert_outputs.view(1, 4, 3, 8) bert_outputs = bert_outputs.sum(2) print(bert_outputs)
gather函数实在是精妙无比，以一段代码解决了复杂的特征融合方式。
参考资料
https://zhuanlan.zhihu.com/p/352877584
https://blog.csdn.net/weixin_42200930/article/details/108995776 (推荐阅读)
