在 mxnet 中高效创建压缩带状对角矩阵

Question

在我的问题中，我有一个包含 n 个元素的向量。给定 window 大小 k 我想有效地创建一个矩阵大小 n x 2k+1 其中包含带状对角线。例如：

a = [a_1, a_2, a_3, a_4]
k = 1
b = [[0, a_1, a_2],
    [a_1, a_2, a_3],
    [a_2, a_3, a_4],
    [a_3, a_4, a_5],
    [a_4, a_5, 0]]

最简单的实现方式是使用 for 循环

out_data = mx.ndarray.zeros((n, 2k+1))
for i in range(0, n):
    for j in range(0, 2k+1):
        index = i - k + j
        if not (index < 0 or index >= seq_len):
            out_data[i][j] = in_data[index]

这很慢。

只需使用 tile 和 reshape 即可轻松创建完整矩阵，但遮罩部分尚不清楚。

更新我找到了一个更快但仍然很慢的实现：

window = 2*self.windowSize + 1
in_data_reshaped = in_data.reshape((batch_size, seq_len))
out_data = mx.ndarray.zeros((seq_len * window))
for i in range(0, seq_len):
    copy_from_start = max(i - self.windowSize, 0)
    copy_from_end = min(seq_len -1, i+1+self.windowSize)
    copy_length = copy_from_end - copy_from_start
    copy_to_start = i*window + (2*self.windowSize + 1 - copy_length)
    copy_to_end = copy_to_start + copy_length
    out_data[copy_to_start:copy_to_end] = in_data_reshaped[copy_from_start:copy_from_end]
out_data = out_data.reshape((seq_len, window))

Answer 1

如果在您的操作中，k 和 n 是常量，您可以使用 mxnet.nd.gather_nd() and mx.nd.scatter_nd 的组合来做您想做的事。尽管生成索引张量同样低效，但因为您只需要执行一次，所以这不是问题。您可能希望使用 gather_nd 有效地 "duplicate" 来自原始数组的数据，然后使用 scatter_nd 将它们分散到最终的矩阵形状。或者，您可以简单地将一个 0 元素连接到您的输入数组（例如 [a_1, a_2, a_3] 会变成 [0, a_1, a_2, a_3]），然后仅使用 mxnet.nd.gather_nd() 将元素复制到您的最终矩阵中.

在 mxnet 中高效创建压缩带状对角矩阵

Efficiently Create compressed banded diagonal matrix in mxnet

python

mxnet