junjue-wang / lovecs Goto Github PK

def replace_bn_with_csn(module:nn.Module, affine_tar=False):
    '''
    Args:
        module: the original model with BNs

    Returns:
        module: the model with the CSNs
    '''
    module_output = module
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        module_output = CSNorm2d(module.num_features, module.eps, module.momentum, module.affine, affine_tar=affine_tar)
        if module.affine:
            with torch.no_grad():
                module_output.weight_source = module.weight#.data.clone().detach()
                module_output.bias_source = module.bias#.data.clone().detach()
        if affine_tar:
            with torch.no_grad():
                module_output.weight_target.data = module.weight#.data.clone().detach()
                module_output.bias_target.data = module.bias#.data.clone().detach()
        module_output.running_mean_source = module.running_mean#.data.clone()
        module_output.running_var_source = module.running_var#.data.clone()
        module_output.running_mean_target = module.running_mean#.data.clone()
        module_output.running_var_target = module.running_var#.data.clone()
        module_output.num_batches_tracked = module.num_batches_tracked#.data.clone()

    for name, child in module.named_children():
        module_output.add_module(name, replace_bn_with_csn(child))
    del module
    return module_output