Pytorch: How to create a random int tensor where a certain percent are of a certain value? For example, 25% are 1s, and rest 0s

user_3259896 0
Viewed 5065

In pytorch I can create a random zero and one tensor with around %50 distribution of each

import torch 
torch.randint(low=0, high=2, size=(2, 5))

I am wondering how I can make a tensor where only 25% of the values are 1s, and the rest are zeros?

pytorch
2 Answers

You can use rand to generate a random tensor between 0,1 and compare that to 0.25:

(torch.rand(size=(2,5)) < 0.25).int()

Output:

tensor([[0, 0, 0, 0, 1],
        [1, 0, 0, 0, 0]], dtype=torch.int32)
user_42384080

Following my answer here: How to randomly set a fixed number of elements in each row of a tensor in PyTorch

Say you want a matrix with dimensions n X d where exactly 25% of the values in each row are 1 and the rest 0, desired_tensor will have the result you want:

n = 2
d = 5
rand_mat = torch.rand(n, d)
k = round(0.25 * d) # For the general case change 0.25 to the percentage you need
k_th_quant = torch.topk(rand_mat, k, largest = False)[0][:,-1:]
bool_tensor = rand_mat <= k_th_quant
desired_tensor = torch.where(bool_tensor,torch.tensor(1),torch.tensor(0))
user_47577150
Related