o

    h                     @   sJ   d dl Z d dlmZ dejdejdejfddZdejdejfdd	ZdS )
    Nconvbnreturnc              	   C   sz   | j }| jdur| jnt|j}|j t|j|j  }t	||
| jdddg | _ t	||j | |j | _| S )zFuse conv and bn into one module.

    Args:
        conv (nn.Module): Conv to be fused.
        bn (nn.Module): BN to be fused.

    Returns:
        nn.Module: Fused module.
    N   )weightbiastorch
zeros_likerunning_meansqrtrunning_varepsnn	Parameterreshapeout_channels)r   r   conv_wconv_bfactor r   6/root/Awesome-Backbones/configs/common/fuse_conv_bn.py
_fuse_conv_bn   s   
r   modulec                 C   s   d}d}|   D ]7\}}t|tjjjtjfr0|du rqt||}|| j|< t	 | j|< d}qt|tj
r;|}|}qt| q| S )a  Recursively fuse conv and bn in a module.

    During inference, the functionary of batch norm layers is turned off
    but only the mean and var alone channels are used, which exposes the
    chance to fuse it with the preceding conv layers to save computations and
    simplify network structures.

    Args:
        module (nn.Module): Module to be fused.

    Returns:
        nn.Module: Fused module.
    N)named_children
isinstancer   modules	batchnorm
_BatchNorm
SyncBatchNormr   _modulesIdentityConv2dfuse_conv_bn)r   Z	last_convZlast_conv_namenamechild
fused_convr   r   r   r"      s"   


r"   )r   torch.nnr   Moduler   r"   r   r   r   r   <module>   s   