coremltools/converters/mil/experimental/passes/generic_conv_batchnorm_fusion.py [71:132]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    is_conv_1d  = len(conv_weight.shape) == 3

    # D_in denotes the spatial dimensions for conv kernel weight
    # for conv_transpose, conv_weight has shape [Cin, Cout / groups, *D_in]
    # for conv, conv_weight has shape [Cout, Cin / groups, *D_in]
    if is_deconv:
        Cout = conv_weight.shape[1] * groups
        Cin = conv_weight.shape[0]
    else:
        Cout = conv_weight.shape[0]
        Cin = conv_weight.shape[1] * groups

    # get the type of the conv weight
    conv_weight_type = conv_weight.dtype

    # create bias for conv if not exist
    if conv_bias is None:
        conv_bias = np.zeros(Cout)
    else:
        conv_bias = conv_bias.val
    conv_bias = conv_bias.astype(conv_weight_type)

    # get the original shape of weight and bias
    origin_weight_shape = conv_weight.shape
    origin_bias_shape = conv_bias.shape

    # update the weight for conv layer
    new_conv_weight = []
    new_conv_bias = []

    if is_deconv:
        conv_weight = np.transpose(conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])
        conv_weight = np.reshape(conv_weight, [Cout, Cin // groups] + list(conv_weight.shape[2:]))

    for i in range(Cout):
        # get batch norm parameters for each channel
        _gamma = gamma[i]
        _beta = beta[i]
        _mean = mean[i]
        _variance = variance[i]
        _scale = _gamma / np.sqrt(_variance + epsilon)

        # get conv weight and bias for each channel
        _conv_weight = conv_weight[i]
        _conv_bias = conv_bias[i]

        # update the conv weight and bias
        _conv_weight = _conv_weight * _scale
        _conv_bias = _scale * (_conv_bias - _mean) + _beta
        new_conv_weight.append(_conv_weight)
        new_conv_bias.append(_conv_bias)

    new_conv_weight = np.array(new_conv_weight).astype(conv_weight_type)
    new_conv_bias = np.array(new_conv_bias).astype(conv_weight_type)

    if is_deconv:
        new_conv_weight = np.reshape(new_conv_weight, [Cout // groups, Cin] + list(new_conv_weight.shape[2:]))
        new_conv_weight = np.transpose(new_conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])

    # make sure the updated weight and bias have the same shape as the original ones
    assert new_conv_weight.shape == origin_weight_shape, "conv weight should have the same shape before and after the fuse_conv_batchnorm pass."
    assert new_conv_bias.shape == origin_bias_shape, "conv bias should have the same shape before and after the fuse_conv_batchnorm pass."
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py [32:93]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    is_conv_1d  = len(conv_weight.shape) == 3

    # D_in denotes the spatial dimensions for conv kernel weight
    # for conv_transpose, conv_weight has shape [Cin, Cout / groups, *D_in]
    # for conv, conv_weight has shape [Cout, Cin / groups, *D_in]
    if is_deconv:
        Cout = conv_weight.shape[1] * groups
        Cin = conv_weight.shape[0]
    else:
        Cout = conv_weight.shape[0]
        Cin = conv_weight.shape[1] * groups

    # get the type of the conv weight
    conv_weight_type = conv_weight.dtype

    # create bias for conv if not exist
    if conv_bias is None:
        conv_bias = np.zeros(Cout)
    else:
        conv_bias = conv_bias.val
    conv_bias = conv_bias.astype(conv_weight_type)

    # get the original shape of weight and bias
    origin_weight_shape = conv_weight.shape
    origin_bias_shape = conv_bias.shape

    # update the weight for conv layer
    new_conv_weight = []
    new_conv_bias = []

    if is_deconv:
        conv_weight = np.transpose(conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])
        conv_weight = np.reshape(conv_weight, [Cout, Cin // groups] + list(conv_weight.shape[2:]))

    for i in range(Cout):
        # get batch norm parameters for each channel
        _gamma = gamma[i]
        _beta = beta[i]
        _mean = mean[i]
        _variance = variance[i]
        _scale = _gamma / np.sqrt(_variance + epsilon)

        # get conv weight and bias for each channel
        _conv_weight = conv_weight[i]
        _conv_bias = conv_bias[i]

        # update the conv weight and bias
        _conv_weight = _conv_weight * _scale
        _conv_bias = _scale * (_conv_bias - _mean) + _beta
        new_conv_weight.append(_conv_weight)
        new_conv_bias.append(_conv_bias)

    new_conv_weight = np.array(new_conv_weight).astype(conv_weight_type)
    new_conv_bias = np.array(new_conv_bias).astype(conv_weight_type)

    if is_deconv:
        new_conv_weight = np.reshape(new_conv_weight, [Cout // groups, Cin] + list(new_conv_weight.shape[2:]))
        new_conv_weight = np.transpose(new_conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])

    # make sure the updated weight and bias have the same shape as the original ones
    assert new_conv_weight.shape == origin_weight_shape, "conv weight should have the same shape before and after the fuse_conv_batchnorm pass."
    assert new_conv_bias.shape == origin_bias_shape, "conv bias should have the same shape before and after the fuse_conv_batchnorm pass."
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -