(self, model)
| 393 | self.backup = {} |
| 394 | |
| 395 | def freeze_bn(self, model): |
| 396 | assert self.backup == {} |
| 397 | for name, m in model.named_modules(): |
| 398 | if isinstance(m, nn.SyncBatchNorm) or isinstance(m, nn.BatchNorm2d): |
| 399 | self.backup[name + '.running_mean'] = m.running_mean.data.clone() |
| 400 | self.backup[name + '.running_var'] = m.running_var.data.clone() |
| 401 | self.backup[name + '.num_batches_tracked'] = m.num_batches_tracked.data.clone() |
| 402 | |
| 403 | def unfreeze_bn(self, model): |
| 404 | for name, m in model.named_modules(): |