如何將 pytorch 張量中維度中的 k 個元素隨機設置為特定值？

Question

例如，如果我有張量（形狀 [2, 3, 5]）

[[[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]]]

並且 k = 2，我想將維度中的 k 個元素（例如，dim=2）隨機設置為特定值（例如，5）：

[[[0.8823, 0.9150, 0.3829, 5, 5],
  [0.6009, 0.2566, 5, 0.9408, 5],
  [5, 5, 0.8694, 0.5677, 0.7411]],

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

Answer 1

使用argsort創建一個隨機二進制掩碼，其中k個True元素的維度為dim並將其設置為 5。

import torch

p = 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]]], requires_grad=True)

k = 2
mask = torch.rand(p.shape).argsort(2) < k
torch.where(mask, 5, p)

Output

tensor([[[5.0000, 0.9150, 5.0000, 0.9593, 0.3904],
         [5.0000, 0.2566, 0.7936, 5.0000, 0.1332],
         [0.9346, 5.0000, 0.8694, 0.5677, 5.0000]],

        [[5.0000, 0.8854, 0.5739, 0.2666, 5.0000],
         [5.0000, 0.4414, 0.2969, 5.0000, 0.1053],
         [5.0000, 5.0000, 0.1994, 0.5472, 0.0062]]], grad_fn=<WhereBackward0>)