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

Question

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

我需要分别针对50个维度在182x91矩阵上运行一个函数。

我需要获取182x91矩阵的对角矩阵条，而我正在使用的功能如下（基于我之前的问题：在numpy或pytorch中自动获取对角矩阵条）：

 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的矩阵，并且不能处理第三维。

所以当我运行这个：

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将 stripe 函数映射到张量的第一维上

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])