How to set minimum k elements in a dimension in pytorch tensor to a specific value?

user_19707146 0
Viewed 47

For example, if I have tensor (shape [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]]]

and k = 2, I want to set minimum k elements in a dimension (e.g., dim=2) to a specific value (e.g., 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]]]
tensor pytorch python
2 Answers

You can extract the k-lowest element per row and mask the initial tensor with that tensor. Given k=2 and v0=5 (the value to replace the k-lowest elements with):

>>> v, _ = x.sort(dim=2)

>>> v[:,:,k:k+1]
tensor([[[0.8823],
         [0.6009],
         [0.7411]],

        [[0.5739],
         [0.2969],
         [0.2695]]])

Doing a slice x[:,:,k:k+1] instead of standard indexing with x[:,:,k] allows to keep the number of dimensions unchanged.

Then we can apply torch.where:

>>> torch.where(x < v[:,:,k:k+1], v0, x)
tensor([[[0.8823, 0.9150, 5.0000, 0.9593, 5.0000],
         [0.6009, 5.0000, 0.7936, 0.9408, 5.0000],
         [0.9346, 5.0000, 0.8694, 5.0000, 0.7411]],

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

Alternatively, you can reassign values on masked-x directly which is in-place:

>>> x[x < v[:,:,k:k+1]] = v0
user_63313690

You can use a combination of torch.topk and torch.Tensor.scatter_.

(Bacause torch.topk return max_top_k and you want min_top_k. We can use -1*all_num for getting min_top_k)

val, ind = torch.topk(-a, k=2)
a.scatter_(index=ind, dim=-1, value=5)
print(a)

tensor([[[0.8823, 0.9150, 5.0000, 0.9593, 5.0000],
         [0.6009, 5.0000, 0.7936, 0.9408, 5.0000],
         [0.9346, 5.0000, 0.8694, 5.0000, 0.7411]],

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

Input:

>>> a = 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]]])


>>> torch.topk(-a, k=2)
# values=tensor(
#     [[[-0.3829, -0.3904],
#       [-0.1332, -0.2566],
#       [-0.5677, -0.5936]],

#      [[-0.2666, -0.4294],
#       [-0.1053, -0.2696],
#       [-0.0062, -0.1994]]]),

# indices=tensor(
#     [[[2, 4],
#       [4, 1],
#       [3, 1]],

#      [[3, 0],
#       [4, 0],
#       [4, 2]]])
user_17405770
Related