从任意尺寸的 PyTorch 张量的特定尺寸/轴获取最终值

Question

假设我有一个 PyTorch 张量，例如：

import torch

x = torch.randn([3, 4, 5])

我想得到一个新的张量，具有相同的维数，包含维度 1 最终值的所有内容。我可以这样做：

x[:, -1:, :]

但是，如果

x

具有任意数量的维度，并且我想从特定维度获取最终值，那么最好的方法是什么？

Answer 1

您可以使用

select

函数（或张量的等效 方法），例如，

dim = 1  # the dimension from which to extract the final values

y = x.select(dim, -1).unsqueeze(dim)

其中

unsqueeze

用于保持与原始张量相同的维数。