Importance-sample timesteps for a batch. :param batch_size: the number of timesteps. :param device: the torch device to save to. :return: a tuple (timesteps, weights): - timesteps: a tensor of timestep indices. - weights: a tensor of
(self, batch_size, device)
| 44 | """ |
| 45 | |
| 46 | def sample(self, batch_size, device): |
| 47 | """ |
| 48 | Importance-sample timesteps for a batch. |
| 49 | :param batch_size: the number of timesteps. |
| 50 | :param device: the torch device to save to. |
| 51 | :return: a tuple (timesteps, weights): |
| 52 | - timesteps: a tensor of timestep indices. |
| 53 | - weights: a tensor of weights to scale the resulting losses. |
| 54 | """ |
| 55 | w = self.weights() |
| 56 | p = w / np.sum(w) |
| 57 | indices_np = np.random.choice(len(p), size=(batch_size, ), p=p) |
| 58 | indices = th.from_numpy(indices_np).long().to(device) |
| 59 | weights_np = 1 / (len(p) * p[indices_np]) |
| 60 | weights = th.from_numpy(weights_np).float().to(device) |
| 61 | return indices, weights |
| 62 | |
| 63 | |
| 64 | class UniformSampler(ScheduleSampler): |