我正在使用HuggingFace Transformers的BERT模型，我希望计算一个句子中各标记的汇总向量（即嵌入），使用mean或max函数。复杂之处在于有些标记是[PAD]，因此在计算平均值或最大值时，我希望忽略这些标记的向量。
这是一个例子。我首先实例化了一个BertTokenizer和一个BertModel：
然后我将一些句子输入到分词器中，得到input_ids和attention_mask。值得注意的是，attention_mask的值为0表示该标记是[PAD]，可以忽略。
sentences = ['Deep learning is difficult yet very rewarding.', 'Deep learning is not easy.', 'But is rewarding if done right.']tokenizer_result = tokenizer(sentences, max_length=32, padding=True, return_attention_mask=True, return_tensors='pt')input_ids = tokenizer_result.input_idsattention_mask = tokenizer_result.attention_maskprint(input_ids.shape) # torch.Size([3, 11])print(input_ids)# tensor([[ 101, 2784, 4083, 2003, 3697, 2664, 2200, 10377, 2075, 1012, 102],# [ 101, 2784, 4083, 2003, 2025, 3733, 1012, 102, 0, 0, 0],# [ 101, 2021, 2003, 10377, 2075, 2065, 2589, 2157, 1012, 102, 0]])print(attention_mask.shape) # torch.Size([3, 11])print(attention_mask)# tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],# [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],# [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]])现在，我调用BERT模型以获取768维的标记嵌入（顶层隐藏状态）。
model_result = model(input_ids, attention_mask=attention_mask, return_dict=True)token_embeddings = model_result.last_hidden_stateprint(token_embeddings.shape) # torch.Size([3, 11, 768])因此，此时我有：
- 标记嵌入在[3, 11, 768]的矩阵中：3个句子，11个标记，每个标记的768维向量。
- 注意力掩码在[3, 11]的矩阵中：3个句子，11个标记。值为1表示非
[PAD]标记。
如何计算有效的、非[PAD]标记的向量的mean / max值？
我尝试使用注意力掩码作为掩码，然后调用torch.max()，但我得到的维度不对：
masked_token_embeddings = token_embeddings[attention_mask==1]print(masked_token_embeddings.shape) # torch.Size([29, 768] <-- 错误的。应该是 [3, 11, 768]pooled = torch.max(masked_token_embeddings, 1)print(pooled.values.shape) # torch.Size([29]) <-- 错误的。应该是 [3, 768]我真正想要的是一个形状为[3, 768]的张量。也就是说，每个句子的768维向量。
回答：
对于max，你可以与attention_mask相乘：
pooled = torch.max((token_embeddings * attention_mask.unsqueeze(-1)), axis=1)对于mean，你可以沿轴求和，然后除以该轴上的attention_mask求和：
mean_pooled = token_embeddings.sum(axis=1) / attention_mask.sum(axis=-1).unsqueeze(-1)