How to sample from a softmax distribution

dist = torch.randn((100, 100))
softmax = nn.Softmax(dim=1)
out = softmax(dist)

This is all pretty standard and makes sense, but I am unable to figure out how to draw random samples from this distribution. In particular, I want to randomly find indices from dist which should be sampled following the softmax distribution.

Is there a way to do this?

You can obtain the probability of sampling for each object by softmax, but you have to have the actual list of objects. Here, I simply assume the list comprises numbers from 0 to 100.

obj_list = list(range(100))
np.random.choice(obj_list, p=dist[0])

note: I used dist[0] since out contains probability distribution for several independent PDFs. You can do similar for the rest.

1 Like

Note that torch.multinomial might be better if you want to stay in PyTorch.
I would recommend to avoid mixing random number generators from different frameworks if you can.

(To me using two RNGs from two different frameworks falls under “being clever” and my rule of thumb here is “don’t try to be clever with random number generation unless you absolutely know what you are doing”… The linked anecdote is not immediately applicable, but it is a trap I have seen people fall into. The other day the big surprise was that PyTorch dataloaders didn’t initialize the numpy seed in workers, that has been fixed in recent PyTorch, but I’d venture that it is a good habit to stay with PyTorch unless you have to use something else.)

Best regards



Awesome, that worked for me, I used torch.multinomial and got the behaviour I was looking for.


1 Like