pytorch按照索引取batch中的数

比如说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))