PyTorch如何将张量值四舍五入到列表中的最近值？

不管用NumPy还是PyTorch，都可以通过广播运算+索引匹配实现高效处理，完全避免逐元素循环，适合大规模张量。

NumPy 实现 #

import numpy as np

# 定义输入张量和目标值列表
input_tensor = np.array([[0.5, 0.25], [1.2, 2]])
target_values = np.array([0, 1, 2])

# 1. 扩展张量维度，实现与目标列表的广播计算绝对差
abs_diffs = np.abs(input_tensor[..., np.newaxis] - target_values)
# 2. 找到每个元素对应最小差的目标值索引
closest_idx = np.argmin(abs_diffs, axis=-1)
# 3. 根据索引映射到最近的目标值
result = target_values[closest_idx]

print(result)
# 输出：
# [[1 0]
#  [1 2]]

PyTorch 实现 #

import torch

input_tensor = torch.tensor([[0.5, 0.25], [1.2, 2]])
target_values = torch.tensor([0, 1, 2])

# 1. 扩展张量维度实现广播
abs_diffs = torch.abs(input_tensor.unsqueeze(-1) - target_values)
# 2. 取最小差对应的索引
closest_idx = torch.argmin(abs_diffs, dim=-1)
# 3. 映射到目标值
result = target_values[closest_idx]

print(result)
# 输出：
# tensor([[1, 0],
#         [1, 2]])

关键说明 #

• 广播机制让我们可以一次性计算所有元素与目标列表的差值，运算效率远高于逐元素循环，尤其适合大尺寸张量。
• 平局处理：示例中0.5刚好在0和1中间，argmin会返回第一个出现最小值的索引（即1）。如果需要自定义平局逻辑（比如优先取较大值），可以在计算差时做调整，例如：

  # NumPy 中优先取较大目标值的平局处理
  # 先计算差，平局时给较大目标值的差减一个极小值，确保它被优先选中
  diffs = input_tensor[..., np.newaxis] - target_values
  abs_diffs = np.abs(diffs)
  # 对差值为0的平局情况，让较大的目标值对应的差更小
  abs_diffs[diffs == 0] -= 1e-9
  closest_idx = np.argmin(abs_diffs, axis=-1)
  

内容的提问来源于stack exchange，提问作者Gooby