| 68 | |
| 69 | |
| 70 | class TimestepSamplerMP: |
| 71 | |
| 72 | def __init__(self, |
| 73 | dp_group: dist.ProcessGroup, |
| 74 | num_train_timesteps: int = 1000) -> None: |
| 75 | self.dp_size = dp_group.size() |
| 76 | self.dp_rank = dp_group.rank() |
| 77 | self.train_timesteps = np.arange(num_train_timesteps) |
| 78 | self.t_step_split = np.array_split(self.train_timesteps, self.dp_size) |
| 79 | self.iter = 0 |
| 80 | |
| 81 | def sample_t(self, |
| 82 | batch_size: int, |
| 83 | training_iter: int = None) -> np.ndarray: |
| 84 | if training_iter is None: |
| 85 | bin_idx = (self.iter + self.dp_rank) % self.dp_size |
| 86 | self.iter += 1 |
| 87 | else: |
| 88 | bin_idx = (training_iter + self.dp_rank) % self.dp_size |
| 89 | bin_steps = self.t_step_split[bin_idx] |
| 90 | return np.random.choice(bin_steps, size=batch_size, replace=False) |
| 91 | |
| 92 | |
| 93 | def get_scheduler(beta_type: str = None, |
nothing calls this directly
no outgoing calls
no test coverage detected