From: https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/
| 2 | |
| 3 | |
| 4 | class AliasMethod(object): |
| 5 | """ |
| 6 | From: https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/ |
| 7 | """ |
| 8 | def __init__(self, probs): |
| 9 | |
| 10 | if probs.sum() > 1: |
| 11 | probs.div_(probs.sum()) |
| 12 | K = len(probs) |
| 13 | self.prob = torch.zeros(K) |
| 14 | self.alias = torch.LongTensor([0]*K) |
| 15 | |
| 16 | # Sort the data into the outcomes with probabilities |
| 17 | # that are larger and smaller than 1/K. |
| 18 | smaller = [] |
| 19 | larger = [] |
| 20 | for kk, prob in enumerate(probs): |
| 21 | self.prob[kk] = K*prob |
| 22 | if self.prob[kk] < 1.0: |
| 23 | smaller.append(kk) |
| 24 | else: |
| 25 | larger.append(kk) |
| 26 | |
| 27 | # Loop though and create little binary mixtures that |
| 28 | # appropriately allocate the larger outcomes over the |
| 29 | # overall uniform mixture. |
| 30 | while len(smaller) > 0 and len(larger) > 0: |
| 31 | small = smaller.pop() |
| 32 | large = larger.pop() |
| 33 | |
| 34 | self.alias[small] = large |
| 35 | self.prob[large] = (self.prob[large] - 1.0) + self.prob[small] |
| 36 | |
| 37 | if self.prob[large] < 1.0: |
| 38 | smaller.append(large) |
| 39 | else: |
| 40 | larger.append(large) |
| 41 | |
| 42 | for last_one in smaller+larger: |
| 43 | self.prob[last_one] = 1 |
| 44 | |
| 45 | def cuda(self): |
| 46 | self.prob = self.prob.cuda() |
| 47 | self.alias = self.alias.cuda() |
| 48 | |
| 49 | def draw(self, N): |
| 50 | """ |
| 51 | Draw N samples from multinomial |
| 52 | :param N: number of samples |
| 53 | :return: samples |
| 54 | """ |
| 55 | K = self.alias.size(0) |
| 56 | |
| 57 | kk = torch.zeros(N, dtype=torch.long, device=self.prob.device).random_(0, K) |
| 58 | prob = self.prob.index_select(0, kk) |
| 59 | alias = self.alias.index_select(0, kk) |
| 60 | # b is whether a random number is greater than q |
| 61 | b = torch.bernoulli(prob) |