Load a checkpoint. Should be called by all processes.
(self, load_dir, step=0)
| 172 | master_params[i].data.copy_(param.data) |
| 173 | |
| 174 | def load(self, load_dir, step=0): |
| 175 | """ |
| 176 | Load a checkpoint. |
| 177 | Should be called by all processes. |
| 178 | """ |
| 179 | if self.is_master: |
| 180 | print(f'\nLoading checkpoint from step {step}...', end='') |
| 181 | |
| 182 | model_ckpts = {} |
| 183 | for name, model in self.models.items(): |
| 184 | model_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')), map_location=self.device, weights_only=True) |
| 185 | model_ckpts[name] = model_ckpt |
| 186 | model.load_state_dict(model_ckpt) |
| 187 | if self.fp16_mode == 'inflat_all': |
| 188 | model.convert_to_fp16() |
| 189 | self._state_dicts_to_master_params(self.master_params, model_ckpts) |
| 190 | del model_ckpts |
| 191 | |
| 192 | if self.is_master: |
| 193 | for i, ema_rate in enumerate(self.ema_rate): |
| 194 | ema_ckpts = {} |
| 195 | for name, model in self.models.items(): |
| 196 | ema_ckpt = torch.load(os.path.join(load_dir, 'ckpts', f'{name}_ema{ema_rate}_step{step:07d}.pt'), map_location=self.device, weights_only=True) |
| 197 | ema_ckpts[name] = ema_ckpt |
| 198 | self._state_dicts_to_master_params(self.ema_params[i], ema_ckpts) |
| 199 | del ema_ckpts |
| 200 | |
| 201 | misc_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'misc_step{step:07d}.pt')), map_location=torch.device('cpu'), weights_only=False) |
| 202 | self.optimizer.load_state_dict(misc_ckpt['optimizer']) |
| 203 | self.step = misc_ckpt['step'] |
| 204 | self.data_sampler.load_state_dict(misc_ckpt['data_sampler']) |
| 205 | if self.fp16_mode == 'amp': |
| 206 | self.scaler.load_state_dict(misc_ckpt['scaler']) |
| 207 | elif self.fp16_mode == 'inflat_all': |
| 208 | self.log_scale = misc_ckpt['log_scale'] |
| 209 | if self.lr_scheduler_config is not None: |
| 210 | self.lr_scheduler.load_state_dict(misc_ckpt['lr_scheduler']) |
| 211 | if self.elastic_controller_config is not None: |
| 212 | self.elastic_controller.load_state_dict(misc_ckpt['elastic_controller']) |
| 213 | if self.grad_clip is not None and not isinstance(self.grad_clip, float): |
| 214 | self.grad_clip.load_state_dict(misc_ckpt['grad_clip']) |
| 215 | del misc_ckpt |
| 216 | |
| 217 | if self.world_size > 1: |
| 218 | dist.barrier() |
| 219 | if self.is_master: |
| 220 | print(' Done.') |
| 221 | |
| 222 | if self.world_size > 1: |
| 223 | self.check_ddp() |
| 224 | |
| 225 | def save(self): |
| 226 | """ |
no test coverage detected