(self, context=None)
| 69 | |
| 70 | @contextmanager |
| 71 | def ema_scope(self, context=None): |
| 72 | if self.use_ema: |
| 73 | self.model_ema.store(self.parameters()) |
| 74 | self.model_ema.copy_to(self) |
| 75 | if context is not None: |
| 76 | print(f"{context}: Switched to EMA weights") |
| 77 | try: |
| 78 | yield None |
| 79 | finally: |
| 80 | if self.use_ema: |
| 81 | self.model_ema.restore(self.parameters()) |
| 82 | if context is not None: |
| 83 | print(f"{context}: Restored training weights") |
| 84 | |
| 85 | def on_train_batch_end(self, *args, **kwargs): |
| 86 | if self.use_ema: |
no test coverage detected