<reponame>hunterkun/small-distiller<filename>swa/utils.py<gh_stars>1-10
import torch

def adjust_learning_rate(optimizer, lr):
    for param_group in optimizer.param_groups:
            param_group['lr'] = lr
    return lr

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

# BatchNorm utils
def _check_bn(module, flag):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        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__, torch.nn.modules.batchnorm._BatchNorm):
        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__, torch.nn.modules.batchnorm._BatchNorm):
        momenta[module] = module.momentum
        
def _set_momenta(module, momenta):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        module.momentum = momenta[module]


def bn_update(loader, model, args):
    if not check_bn(model):
        return
    was_training=model.training
    model.train()
    momenta = {}
    model.apply(reset_bn)
    model.apply(lambda module: _get_momenta(module, momenta))
    n = 0
    for input in loader:
        if isinstance(input, (list, tuple)):
            input = input[0]
        b = input.size(0)
        momentum = b / float(n + b)
        for module in momenta.keys():
            module.momentum = momentum
        input = input.to(args.device)
        model(input)
        n += b
    model.apply(lambda module: _set_momenta(module, momenta))
    model.train(was_training)
import torch

def adjust_learning_rate(optimizer, lr):
    for param_group in optimizer.param_groups:
            param_group['lr'] = lr
    return lr

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

# BatchNorm utils
def _check_bn(module, flag):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        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__, torch.nn.modules.batchnorm._BatchNorm):
        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__, torch.nn.modules.batchnorm._BatchNorm):
        momenta[module] = module.momentum
        
def _set_momenta(module, momenta):
    if issubclass(module.__class__, torch.nn.modules.batchnorm._BatchNorm):
        module.momentum = momenta[module]


def bn_update(loader, model, args):
    if not check_bn(model):
        return
    was_training=model.training
    model.train()
    momenta = {}
    model.apply(reset_bn)
    model.apply(lambda module: _get_momenta(module, momenta))
    n = 0
    for input in loader:
        if isinstance(input, (list, tuple)):
            input = input[0]
        b = input.size(0)
        momentum = b / float(n + b)
        for module in momenta.keys():
            module.momentum = momentum
        input = input.to(args.device)
        model(input)
        n += b
    model.apply(lambda module: _set_momenta(module, momenta))
    model.train(was_training)
# BatchNorm utils