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

问题描述

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

I have a tensor of size [150, 182, 91], the first part is just the batch size while the matrix I am interested in is the 182x91 one.

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

I need to run a function on the 182x91 matrix for each of the 50 dimensions separately.

我需要得到一个 182x91 矩阵的对角矩阵条纹，我使用的函数如下(基于我之前的问题:在 numpy 或 pytorch 中自动获取对角矩阵条纹):

I need to get a diagonal matrix stripe of the 182x91 matrix, and the function I am using is the following one (based on my previous question: Getting diagonal matrix stripe automatically in numpy or 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 的矩阵，无法处理第 3 维.

The stripe function expects a matrix of size IxJ and can't deal with the 3rd dimension.

所以当我运行这个时:

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

我收到此错误:ValueError:解包的值太多(预期为 2)

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

If I just try to skip the first dimension by doing x, i, j = a.size(),I get this: RuntimeError: invalid argument 1: expected a matrix or a vector at

我仍在使用 PyTorch 0.3.1.任何帮助表示赞赏！

I am still on PyTorch 0.3.1. Any help is appreciated!

推荐答案

您可以使用 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])

这篇关于如何将自定义函数应用于 PyTorch 矩阵中的特定列的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！