pytorch：如何在 4-d 张量的所有单元上应用 function

Question

我正在尝试在具有以下维度的 4-D 张量上应用 function（我认为它是一个二维矩阵，每个单元格中都有一个二维矩阵）：[N x N x N x N]。 应用 function 返回 [1 x N] 张量，因此在应用 function 之后，我期待以下尺寸的张量：[N x N x 1 x N]。

示例：让我们定义 [4 x 4 x 4 x 4] 张量：

tensor_4d = torch.tensor([[[[1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[1., 0., 0., 0.],
  [1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[1., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]]],
[[[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 1., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.]]],
[[[0., 0., 1., 0.],
  [0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 0.]]],
[[[0., 0., 1., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 1.],
  [0., 0., 0., 1.]],
 [[0., 0., 0., 0.],
  [0., 0., 1., 0.],
  [0., 0., 0., 1.],
  [0., 0., 0., 1.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 1.],
  [0., 0., 0., 1.]],
 [[0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 0.],
  [0., 0., 0., 1.]]]], dtype=torch.float64)

让我们看看 [3][0] 处的 tensor_4d：

tensor_4d[3][0]
tensor([[0., 0., 1., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 1.],
        [0., 0., 0., 1.]], dtype=torch.float64)

这是我的申请 function：

def apply_function(tensor_2d):
    eigenvalues, eigenvectors = torch.eig(input=tensor_2d, eigenvectors=True)
    return eigenvector[:, 2]

这是应用 function 的结果：

apply_function(tensor_4d[3][0])
tensor([-1.0000e+00,  0.0000e+00, 4.0083e-292,  0.0000e+00],
       dtype=torch.float64)

所以 apply_function 适用于每个单元格。 接下来，我尝试将 apply_function 与整个矩阵一起使用，并期望每个单元格都包含为此单元格激活“apply_function”的结果。 但是，当使用应用 function 时，我收到以下错误：

apply_function(tensor_4d)

Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2019.2.3\helpers\pydev\_pydevd_bundle\pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "C:/Users/LiavB/PycharmProjects/RBC/Components/RBC_torch.py", line 41, in apply_function
    eigenvalues, eigenvectors = torch.eig(input=tensor_2d, eigenvectors=True)
RuntimeError: invalid argument 1: A should be 2 dimensional at ..\aten\src\TH/generic/THTensorLapack.cpp:206

Answer 1

我们试试看：

new_shape=(-1,)+tensor_4d.shape[2:]

out = (torch.stack([apply_function(t) for t in tensor_4d.view(new_shape)], axis=-1)
      .reshape(new_shape)
)