如何在PyTorch张量的指定维度中将最小k个元素设为特定值？

假设我们有一个形状为[2, 3, 5]的张量：

import torch

input_tensor = torch.tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904],
                             [0.6009, 0.2566, 0.7936, 0.9408, 0.1332],
                             [0.9346, 0.5936, 0.8694, 0.5677, 0.7411]],

                            [[0.4294, 0.8854, 0.5739, 0.2666, 0.6274],
                             [0.2696, 0.4414, 0.2969, 0.8317, 0.1053],
                             [0.2695, 0.3588, 0.1994, 0.5472, 0.0062]]])

需要实现的是：在指定维度（如dim=2）中，将最小的k=2个元素替换为特定值（如5），最终得到结果：

[[[0.8823, 0.9150, 5, 0.9593, 5],
  [0.6009, 5, 0.7936, 0.9408, 5],
  [0.9346, 5, 0.8694, 5, 0.7411]],

 [[5, 0.8854, 0.5739, 5, 0.6274],
  [5, 0.4414, 0.2969, 0.8317, 5],
  [0.2695, 0.3588, 5, 0.5472, 5]]]

实现方法（基于PyTorch） #

通过以下步骤完成需求：

1. 获取指定维度上第k小的元素阈值
2. 创建布尔掩码标记需要替换的元素位置
3. 克隆原张量并替换目标位置为指定值

代码实现：

def replace_smallest_k_elements(tensor, dim, k, target_value):
    # 获取指定维度第k小的元素值，keepdim保证维度匹配
    kth_value, _ = torch.kthvalue(tensor, k, dim=dim, keepdim=True)
    # 生成掩码：标记出小于等于第k小值的元素
    mask = tensor <= kth_value
    # 克隆原张量避免修改原始数据，替换对应位置
    result = tensor.clone()
    result[mask] = target_value
    return result

# 调用函数得到结果
output_tensor = replace_smallest_k_elements(input_tensor, dim=2, k=2, target_value=5)
print(output_tensor)

说明 #

• torch.kthvalue精准获取指定维度的第k小元素，keepdim=True确保返回形状与原张量匹配，方便掩码对齐。
• 布尔掩码tensor <= kth_value直接定位所有最小的k个元素位置。
• 使用clone()复制原张量，避免操作过程中修改输入数据。

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