| 338 | |
| 339 | |
| 340 | class EMA(): |
| 341 | def __init__(self, beta): |
| 342 | super().__init__() |
| 343 | self.beta = beta |
| 344 | |
| 345 | def update_model_average(self, ma_model, current_model): |
| 346 | for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()): |
| 347 | old_weight, up_weight = ma_params.data, current_params.data |
| 348 | ma_params.data = self.update_average(old_weight, up_weight) |
| 349 | |
| 350 | def update_average(self, old, new): |
| 351 | if old is None: |
| 352 | return new |
| 353 | return old * self.beta + (1 - self.beta) * new |
| 354 | |
| 355 | |
| 356 | def get_pad(in_, ksize, stride, atrous=1): |