|
比如说bert的输出表征是基于子词的，想要用于基于词的任务，需要将词对应的最后一个子词的表征取出来，代码如下：
import torch
import torch.nn as nn
torch.manual_seed(1)
mix = torch.randn([2, 5, 3])
range_vector = torch.tensor([[0], [1]])
offsets2d = torch.tensor([[1, 3, 0], [1, 2, 4]])
print(mix)
selected_embeddings = mix[range_vector, offsets2d]
print(selected_embeddings.size())
print(selected_embeddings)
结果：
tensor([[[-1.5256, -0.7502, -0.6540],
[-1.6095, -0.1002, -0.6092],
[-0.9798, -1.6091, -0.7121],
[ 0.3037, -0.7773, -0.2515],
[-0.2223, 1.6871, -0.3206]],
[[-0.2993, 1.8793, -0.0721],
[ 0.1578, -0.7735, 0.1991],
[ 0.0457, -1.3924, 2.6891],
[-0.1110, 0.2927, -0.1578],
[-0.0288, 2.3571, -1.0373]]])
torch.Size([2, 3, 3])
tensor([[[-1.6095, -0.1002, -0.6092],
[ 0.3037, -0.7773, -0.2515],
[-1.5256, -0.7502, -0.6540]],
[[ 0.1578, -0.7735, 0.1991],
[ 0.0457, -1.3924, 2.6891],
[-0.0288, 2.3571, -1.0373]]])
Process finished with exit code 0
此段代码，在模型转成onnx时候会报错，可改成
mix = torch.randn([2, 5, 3]).cuda()
offsets = torch.tensor([[1, 3, 0], [1, 2, 4]]).cuda()
# 按索引取数
B, S, D = mix.size()
new_mix = mix.view(-1, D)
_, W = offsets.size()
right_add = torch.arange(0, B).unsqueeze(-1).cuda()
right_add = right_add * S
right_add.expand([B, W])
new_offsets = right_add + offsets
new_offsets = new_offsets.view(-1)
print(new_offsets)
out1 = new_mix.index_select(0, new_offsets)
# index_select 必须是一维向量
# torch.gather输出维度和输入的维度必须相同
print(out1.view(B, W, -1))
|