
from __future__ import absolute_import, division, print_function, unicode_literals

import copy
import torch

def fuse_conv_bn_eval(conv, bn):
    assert(not (conv.training or bn.training)), "Fusion only for eval!"
    fused_conv = copy.deepcopy(conv)

    w_conv = fused_conv.weight
    b_conv = fused_conv.bias

    bn_mean = bn.running_mean
    bn_var_sqrt = torch.sqrt(bn.running_var + bn.eps)

    bn_weight = bn.weight
    bn_bias = bn.bias

    if b_conv is None:
        b_conv = bn_mean.new_zeros(bn_mean.shape)

    w_conv = w_conv * (bn_weight / bn_var_sqrt).reshape([-1, 1, 1, 1])
    b_conv = (b_conv - bn_mean) / bn_var_sqrt * bn_weight + bn_bias

    fused_conv.weight = torch.nn.Parameter(w_conv)
    fused_conv.bias = torch.nn.Parameter(b_conv)

    return fused_conv
