如何将自定义函数应用于 PyTorch 矩阵中的特定列

Question

我有一个大小为 [150, 182, 91] 的张量，第一部分只是批量大小，而我感兴趣的矩阵是 182x91。

我需要运行在 182x91 矩阵上分别为 50 个维度中的每个维度创建一个函数。

我需要得到182x91矩阵的对角矩阵条纹，我使用的函数是下面这个（根据我之前的问题：）：

 def stripe(a):

    i, j = a.size()
    assert (i >= j)

    out = torch.zeros((i - j + 1, j))
    for diag in range(0, i - j + 1):
        out[diag] = torch.diag(a, -diag)
    return out

stripe 函数需要大小为 IxJ 的矩阵，无法处理第 3 维。

所以当我运行这个：

some_matrix = x # <class 'torch.autograd.variable.Variable'> torch.Size([150, 182, 91])
get_diag = stripe(some_matrix)

我得到这个错误：ValueError: too many values to unpack (expected 2)

如果我只是尝试通过 x, i, j = a.size() 跳过第一个维度，我明白了：RuntimeError: invalid argument 1: expected a matrix or a vector at

我还在使用 PyTorch 0.3.1。感谢您的帮助！

Answer 1

您可以使用 torch.unbind as

将条纹函数映射到张量的第一维

In [1]: import torch

In [2]: def strip(a):
   ...:     i, j = a.size()
   ...:     assert(i >= j)
   ...:     out = torch.zeros((i - j + 1, j))
   ...:     for diag in range(0, i - j + 1):
   ...:         out[diag] = torch.diag(a, -diag)
   ...:     return out
   ...: 
   ...: 

In [3]: a = torch.randn((182, 91)).cuda()

In [5]: output = strip(a)

In [6]: output.size()
Out[6]: torch.Size([92, 91])

In [7]: a = torch.randn((150, 182, 91))

In [8]: output = list(map(strip, torch.unbind(a, 0)))

In [9]: output = torch.stack(output, 0)

In [10]: output.size()
Out[10]: torch.Size([150, 92, 91])

Answer 2

这是一种不使用 stack 和 unbind 的方法，通过直接在批处理矩阵上计算对角线条纹：

def batch_stripe(a):
    b, i, j = a.size()
    assert i > j
    b_s, k, l = a.stride()
    return torch.as_strided(a, (b, i - j, j), (b_s, k, k+1))

更多请参考： https://discuss.pytorch.org/t/optimizing-diagonal-stripe-code/17777/5

如何将自定义函数应用于 PyTorch 矩阵中的特定列

How to apply a custom function to specific columns in a matrix in PyTorch

python

matrix

neural-network

pytorch

tensor