Pytorch，在将一个向量的排列添加到矩阵的条目时摆脱 for 循环？

Question

我正在尝试实现这篇论文，并坚持这个简单的步骤。虽然这与注意力有关，但我坚持的是如何在不使用 for 循环的情况下实现添加到矩阵的向量的排列。

注意力分数添加了学习偏差向量，理论上它编码分数代表的两个标记的相对位置 (j-i)

所以 alpha 是一个 T x T 矩阵，T 取决于被转发的 batch，B 是一个学习的偏置向量，它的长度必须是固定的并且大到 2T。我认为我目前的实现符合论文的建议：

    def __init__(...):
       ...
        self.bias = torch.nn.Parameter(torch.randn(config.n),requires_grad = True)
        stdv = 1. / math.sqrt(self.bias.data.size(0))
        self.bias.data.uniform_(-stdv, stdv)
     def forward(..)
        ...
        #n = 201  (2* max_seq_len + 1)

        B_matrix = torch.zeros(self.T, self.T) # 60 x 60
        for i in range(self.T):
          B_matrix[i] = self.bias[torch.arange(start=n//2-i, end=n//2-i+T)])]

        attention_scores = attention_scores + B_matrix.unsqueeze(0)
        # 64 x 60 x 60   
        ...

这是唯一相关的部分

B_matrix = torch.zeros(self.T, self.T) # 60 x 60
        for i in range(self.T):
          B_matrix[i] = self.bias[torch.arange(start=n//2-i, end=n//2-i+T)])]

基本上尽量不使用 for 循环遍历每一行。

但我知道这一定非常低效，而且当这个模型非常大时成本很高。我正在对每一行进行显式 for 循环以获得学习偏差向量的排列。

任何人都可以通过智能广播帮助我找到更好的方法吗？

想了想，不用实例化一个零矩阵，但还是摆脱不了for循环？并且不能使用收集，因为 B_matrix 与平铺的 b 向量大小不同。

functor = lambda i : bias[torch.arange(start=n//2-i, end=n//2-i+T)]
B_matrix = torch.stack([functor(i) for i in torch.arange(T)])

Answer 1

我无法弄清楚 n 应该在您的代码中是什么，但我认为以下使用 torch.meshgrid 的示例提供了您正在寻找的内容。

假设

n, m = 10, 20   # arbitrary
a = torch.randn(n, m)
b = torch.randn(n + m)

然后

for i in range(n):
    for j in range(m):
        a[i, j] = a[i, j] + b[n - i + j]

等同于

ii, jj = torch.meshgrid(torch.arange(n), torch.arange(m))
a = a + b[n - ii + jj]

尽管后者是一个 out-of-place 操作，这通常是一件好事。如果您确实想要 in-place 操作，请将 a = 替换为 a[...] =。

请注意，这是 integer array indexing 的示例，其中我们使用与 a.

形状相同的张量对 b 进行索引

Pytorch，在将一个向量的排列添加到矩阵的条目时摆脱 for 循环？

Pytorch, get rid of a for loop when adding permutation of one vector to entries of a matrix?

for-loop

numpy

matrix

attention-model

pytorch