A distribution over timesteps in the diffusion process, intended to reduce variance of the objective. By default, samplers perform unbiased importance sampling, in which the objective's mean is unchanged. However, subclasses may override sample() to change how the resampled
| 27 | |
| 28 | |
| 29 | class ScheduleSampler(ABC): |
| 30 | """ |
| 31 | A distribution over timesteps in the diffusion process, intended to reduce |
| 32 | variance of the objective. |
| 33 | By default, samplers perform unbiased importance sampling, in which the |
| 34 | objective's mean is unchanged. |
| 35 | However, subclasses may override sample() to change how the resampled |
| 36 | terms are reweighted, allowing for actual changes in the objective. |
| 37 | """ |
| 38 | |
| 39 | @abstractmethod |
| 40 | def weights(self): |
| 41 | """ |
| 42 | Get a numpy array of weights, one per diffusion step. |
| 43 | The weights needn't be normalized, but must be positive. |
| 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): |
nothing calls this directly
no outgoing calls
no test coverage detected