(self, ma_model, current_model)
| 196 | |
| 197 | @torch.no_grad() |
| 198 | def update_moving_average(self, ma_model, current_model): |
| 199 | copy, lerp = self.inplace_copy, self.inplace_lerp |
| 200 | current_decay = self.get_current_decay() |
| 201 | |
| 202 | for (name, current_params), (_, ma_params) in zip(self.get_params_iter(current_model), self.get_params_iter(ma_model)): |
| 203 | if name in self.ignore_names: |
| 204 | continue |
| 205 | |
| 206 | if any([name.startswith(prefix) for prefix in self.ignore_startswith_names]): |
| 207 | continue |
| 208 | |
| 209 | if name in self.param_or_buffer_names_no_ema: |
| 210 | copy(ma_params.data, current_params.data) |
| 211 | continue |
| 212 | |
| 213 | lerp(ma_params.data, current_params.data, 1. - current_decay) |
| 214 | |
| 215 | for (name, current_buffer), (_, ma_buffer) in zip(self.get_buffers_iter(current_model), self.get_buffers_iter(ma_model)): |
| 216 | if name in self.ignore_names: |
| 217 | continue |
| 218 | |
| 219 | if any([name.startswith(prefix) for prefix in self.ignore_startswith_names]): |
| 220 | continue |
| 221 | |
| 222 | if name in self.param_or_buffer_names_no_ema: |
| 223 | copy(ma_buffer.data, current_buffer.data) |
| 224 | continue |
| 225 | |
| 226 | lerp(ma_buffer.data, current_buffer.data, 1. - current_decay) |
| 227 | |
| 228 | def __call__(self, *args, **kwargs): |
| 229 | return self.ema_model(*args, **kwargs) |
no test coverage detected