|
|
|
|
|
|
| """
|
| @Author : Peike Li
|
| @Contact : peike.li@yahoo.com
|
| @File : schp.py
|
| @Time : 4/8/19 2:11 PM
|
| @Desc :
|
| @License : This source code is licensed under the license found in the
|
| LICENSE file in the root directory of this source tree.
|
| """
|
|
|
| import os
|
| import torch
|
| import modules
|
|
|
| def moving_average(net1, net2, alpha=1):
|
| for param1, param2 in zip(net1.parameters(), net2.parameters()):
|
| param1.data *= (1.0 - alpha)
|
| param1.data += param2.data * alpha
|
|
|
|
|
| def _check_bn(module, flag):
|
| if issubclass(module.__class__, modules.bn.InPlaceABNSync):
|
| flag[0] = True
|
|
|
|
|
| def check_bn(model):
|
| flag = [False]
|
| model.apply(lambda module: _check_bn(module, flag))
|
| return flag[0]
|
|
|
|
|
| def reset_bn(module):
|
| if issubclass(module.__class__, modules.bn.InPlaceABNSync):
|
| module.running_mean = torch.zeros_like(module.running_mean)
|
| module.running_var = torch.ones_like(module.running_var)
|
|
|
|
|
| def _get_momenta(module, momenta):
|
| if issubclass(module.__class__, modules.bn.InPlaceABNSync):
|
| momenta[module] = module.momentum
|
|
|
|
|
| def _set_momenta(module, momenta):
|
| if issubclass(module.__class__, modules.bn.InPlaceABNSync):
|
| module.momentum = momenta[module]
|
|
|
|
|
| def bn_re_estimate(loader, model):
|
| if not check_bn(model):
|
| print('No batch norm layer detected')
|
| return
|
| model.train()
|
| momenta = {}
|
| model.apply(reset_bn)
|
| model.apply(lambda module: _get_momenta(module, momenta))
|
| n = 0
|
| for i_iter, batch in enumerate(loader):
|
| images, labels, _ = batch
|
| b = images.data.size(0)
|
| momentum = b / (n + b)
|
| for module in momenta.keys():
|
| module.momentum = momentum
|
| model(images)
|
| n += b
|
| model.apply(lambda module: _set_momenta(module, momenta))
|
|
|
|
|
| def save_schp_checkpoint(states, is_best_parsing, output_dir, filename='schp_checkpoint.pth.tar'):
|
| save_path = os.path.join(output_dir, filename)
|
| if os.path.exists(save_path):
|
| os.remove(save_path)
|
| torch.save(states, save_path)
|
| if is_best_parsing and 'state_dict' in states:
|
| best_save_path = os.path.join(output_dir, 'model_parsing_best.pth.tar')
|
| if os.path.exists(best_save_path):
|
| os.remove(best_save_path)
|
| torch.save(states, best_save_path)
|
|
|
