Store a copy of the EMA params in fp32. If state dict is passed, the EMA params is copied from the provided state dict. Otherwise, it is copied from the current EMA model parameters.
(self, state_dict=None)
| 43 | self.update_freq_counter = 0 |
| 44 | |
| 45 | def build_fp32_params(self, state_dict=None): |
| 46 | """ |
| 47 | Store a copy of the EMA params in fp32. |
| 48 | If state dict is passed, the EMA params is copied from |
| 49 | the provided state dict. Otherwise, it is copied from the |
| 50 | current EMA model parameters. |
| 51 | """ |
| 52 | if not self.ema_fp32: |
| 53 | raise RuntimeError( |
| 54 | "build_fp32_params should not be called if ema_fp32=False. " |
| 55 | "Use ema_fp32=True if this is really intended." |
| 56 | ) |
| 57 | |
| 58 | if state_dict is None: |
| 59 | state_dict = self.model.state_dict() |
| 60 | |
| 61 | def _to_float(t): |
| 62 | """Internal: to float. |
| 63 | |
| 64 | Args: |
| 65 | t: TODO. |
| 66 | """ |
| 67 | return t.float() if torch.is_floating_point(t) else t |
| 68 | |
| 69 | for param_key in state_dict: |
| 70 | if param_key in self.fp32_params: |
| 71 | self.fp32_params[param_key].copy_(state_dict[param_key]) |
| 72 | else: |
| 73 | self.fp32_params[param_key] = _to_float(state_dict[param_key]) |
| 74 | |
| 75 | def restore(self, state_dict, build_fp32_params=False): |
| 76 | """Load data from a model spec into EMA model""" |
no test coverage detected