如何在PyTorch中将自定义函数应用于矩阵中的特定列 [英] How to apply a custom function to specific columns in a matrix in 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.

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

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

我需要获取182x91矩阵的对角矩阵条纹，并且我正在使用的功能如下(基于我之前的问题:

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.

因此，当我运行此代码时:

So when I run this:

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

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中将自定义函数应用于矩阵中的特定列的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！