tensor用方括号取值 给定坐标 选取部分元素

一、 问题描述

已知BATCH_SIZE个坐标，从n*n的矩阵中取出BATCH_SIZE个值。

二、 问题背景

上周自己写CRF的时候遇到了一个问题，

在pytorch官方的BI-LSTM CRF教程中，代码如下

def _score_sentence(self, feats, tags):
        # Gives the score of a provided tag sequence
        score = torch.zeros(1)
        tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])
        for i, feat in enumerate(feats):
            score = score + \
                self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
        score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
        return score

输入向量没有batch_size维度，输入是一句话对应的特征和标签，输出是一个数。

我就想在实现的时候加一个batch_size维度就好吧，结果代码写成了这样：

    def _score_sentence(self, feats, tags):
        # Gives the score of a provided tag sequence
        score = torch.zeros(BATCH_SIZE, 1).to(device)
        temp_tensor = torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long).to(device)
        temp_tensor = temp_tensor.expand(BATCH_SIZE, 1)
        tags = torch.cat([temp_tensor, tags], dim=-1)
        for i, feat in enumerate(feats.permute([1, 0, 2])):
            for b in range(BATCH_SIZE):
                tran = self.transitions[tags[b, i + 1], tags[b, i]]
                score[b] = score[b] + tran + feat[b, tags[b, i + 1]]
        for b in range(BATCH_SIZE):
            score[b] = score[b] + self.transitions[self.tag_to_ix[STOP_TAG], tags[b, -1]]
        return score

迭代计算batch中每一个句子的得分，然后赋值给score的每一位。
在实际运行过程中，计算的速度特别慢，可能相当于放弃了并行化计算框架的优势。
所以就想到能不能有方法实现这种功能：从n*n的转移矩阵中取出BATCH_SIZE个值。

在pytorch官方文档上找了很多相关张量操作，比如index_select、masked_select、scatter、gather，然而都没有完美解决问题。

三、 解决方案

在自己写的CRF太慢之后，就用了torchcrf这个包。
今天看了这个包里的代码，这部分是这样的：

		# emissions: (seq_length, batch_size, num_tags)
        # tags: (seq_length, batch_size)
        # mask: (seq_length, batch_size)
        for i in range(1, seq_length):
            # Transition score to next tag, only added if next timestep is valid (mask == 1)
            # shape: (batch_size,)
            score += self.transitions[tags[i - 1], tags[i]] * mask[i]

            # Emission score for next tag, only added if next timestep is valid (mask == 1)
            # shape: (batch_size,)
            score += emissions[i, torch.arange(batch_size), tags[i]] * mask[i]

经过总结得知：

M[x, y]

1. 当M是n*n的矩阵，x, y是合法的整数时，能取到矩阵里的一个值；
2. 当x, y是分号隔开的整数时，返回一个子矩阵；
3. 当x, y都是长度为b的整数向量时，返回一个长度为b的向量[M[x[0], y[0]], M[x[1], y[1]],……,M[x[b-1], y[b-1]]]
4. 当x是一个整数, y是长度为b的整数向量时，返回一个长度为b的向量[M[x, y[0]], M[x, y[1]],……,M[x, y[b-1]]]，反之亦然。

异常情况：当x, y为长度不等的向量时

import torch
t = torch.arange(0, 16).reshape(4, -1)
print(t)

t1 = t[[1, 3], [0, 1, 2]]
print(t1)
print(t.shape, t1.shape)

会有如下异常：

IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [2], [3]

四、 示例

上图表示取坐标为[1, 0], [3, 1], [1,2]的三个元素。

本文只选取了被选矩阵为二阶张量（二维矩阵）的情况作为举例。
对于更高阶张量的取值，只需要在方括号中对应位置上添加索引即可。

M[x, y, z]

五、 补充

M[mask]

当mask是一个bool类型的矩阵，且与M形状相同时，返回一个新的一维张量，该张量mask作为掩码对M进行索引。

import torch
t = torch.arange(0, 16).reshape(4, -1)
print(t)

mask = t > 7
print(mask)
print(mask.dtype)

t1 = t[mask]
print(t1)
print(t1.dtype)

tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
tensor([[False, False, False, False],
        [False, False, False, False],
        [ True,  True,  True,  True],
        [ True,  True,  True,  True]])
torch.bool
tensor([ 8,  9, 10, 11, 12, 13, 14, 15])
torch.int64