从任意维度PyTorch张量中提取指定维度的最终值-Python教程-PHP中文网

从任意维度PyTorch张量中提取指定维度的最终值

碧海醫心

发布： 2025-10-24 10:35:01

原创

731人浏览过

从任意维度pytorch张量中提取指定维度的最终值

本文介绍了如何从任意维度的PyTorch张量中提取特定维度的最后一个值，并保持张量的维度不变。主要利用 `torch.index_select` 函数选择指定维度的最后一个索引，并通过 `squeeze` 函数去除不必要的维度，从而获得目标张量。本文提供了详细的代码示例和使用说明，帮助读者理解和应用该方法。

在处理PyTorch张量时，经常需要提取特定维度的信息。当需要获取某个维度的最后一个值时，torch.index_select 函数提供了一种灵活且通用的解决方案。本文将详细介绍如何使用 torch.index_select 从任意维度的PyTorch张量中提取指定维度的最终值，并讨论如何根据需要调整结果张量的维度。

使用 torch.index_select 提取最终值

torch.index_select(input, dim, index) 函数允许我们沿着指定的维度 dim，根据 index 提取张量 input 的元素。为了提取指定维度的最后一个值，我们可以将 index 设置为该维度的最后一个索引。

以下代码展示了如何使用 torch.index_select 提取张量 x 的维度 dim 的最后一个值：

import torch

def get_last_value(x, dim):
  """
  从张量 x 的指定维度 dim 中提取最后一个值。

  Args:
    x: 输入张量。
    dim: 要提取最后一个值的维度。

  Returns:
    一个与输入张量具有相同维度的张量，其中指定维度仅包含最后一个值。
  """
  return torch.index_select(x, dim=dim, index=torch.tensor(x.size(dim) - 1))

# 示例
x = torch.randn([3, 4, 5])
dim = 1
result = get_last_value(x, dim)
print(f"原始张量形状: {x.shape}")
print(f"提取后的张量形状: {result.shape}")

登录后复制

在上述代码中，torch.index_select 函数返回一个新的张量，该张量与原始张量 x 具有相同的维度，但在指定的维度 dim 上，它只包含最后一个值。例如，如果 x 的形状是 [3, 4, 5]，并且 dim 是 1，那么 result 的形状将是 [3, 1, 5]。

百度文心百中

百度大模型语义搜索体验中心

查看详情

使用 squeeze 函数去除多余维度

有时，我们可能希望去除提取后张量中维度为 1 的维度。例如，在上面的例子中，我们可能希望将 result 的形状从 [3, 1, 5] 变为 [3, 5]。这时，可以使用 squeeze 函数。

以下代码展示了如何使用 squeeze 函数去除多余维度：

import torch

def get_last_value_and_squeeze(x, dim):
  """
  从张量 x 的指定维度 dim 中提取最后一个值，并去除该维度。

  Args:
    x: 输入张量。
    dim: 要提取最后一个值的维度。

  Returns:
    一个张量，其中指定维度的最后一个值被提取，并且该维度已被去除。
  """
  return torch.index_select(x, dim=dim, index=torch.tensor(x.size(dim) - 1)).squeeze(dim=dim)

# 示例
x = torch.randn([3, 4, 5])
dim = 1
result = get_last_value_and_squeeze(x, dim)
print(f"原始张量形状: {x.shape}")
print(f"提取并去除维度后的张量形状: {result.shape}")

登录后复制

在这个例子中，squeeze(dim=dim) 函数会去除 result 中维度为 dim 的维度，从而将 result 的形状从 [3, 1, 5] 变为 [3, 5]。