| 70 | |
| 71 | @contextmanager |
| 72 | def ema_scope(self, context=None): |
| 73 | if self.use_ema: |
| 74 | self.model_ema.store(self.parameters()) |
| 75 | self.model_ema.copy_to(self) |
| 76 | if context is not None: |
| 77 | print(f"{context}: Switched to EMA weights") |
| 78 | try: |
| 79 | yield None |
| 80 | finally: |
| 81 | if self.use_ema: |
| 82 | self.model_ema.restore(self.parameters()) |
| 83 | if context is not None: |
| 84 | print(f"{context}: Restored training weights") |
| 85 | |
| 86 | def init_from_ckpt(self, path, ignore_keys=list()): |
| 87 | sd = torch.load(path, map_location="cpu")["state_dict"] |