(self, rate)
| 127 | dist_util.sync_params(self.model.parameters()) |
| 128 | |
| 129 | def _load_ema_parameters(self, rate): |
| 130 | ema_params = copy.deepcopy(self.master_params) |
| 131 | |
| 132 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint |
| 133 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) |
| 134 | if ema_checkpoint: |
| 135 | if dist.get_rank() == 0: |
| 136 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") |
| 137 | state_dict = dist_util.load_state_dict( |
| 138 | ema_checkpoint, map_location=dist_util.dev() |
| 139 | ) |
| 140 | ema_params = self._state_dict_to_master_params(state_dict) |
| 141 | |
| 142 | dist_util.sync_params(ema_params) |
| 143 | return ema_params |
| 144 | |
| 145 | def _load_optimizer_state(self): |
| 146 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint |
no test coverage detected