o

    h                     @   sR   d dl mZ d dlmZmZmZmZ d dlmZm	Z	 d dl
Z
G dd dejZdS )    N)build_conv_layerbuild_norm_layerbuild_activation_layerbuild_padding_layer)kaiming_init
constant_initc                       s^   e Zd ZdZdddddddedddd	d
df fdd
	Zedd Zdd ZdddZ	  Z
S )
ConvModulea
  A conv block that bundles conv/norm/activation layers.

    This block simplifies the usage of convolution layers, which are commonly
    used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
    It is based upon three build methods: `build_conv_layer()`,
    `build_norm_layer()` and `build_activation_layer()`.

    Besides, we add some additional features in this module.
    1. Automatically set `bias` of the conv layer.
    2. Spectral norm is supported.
    3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only
    supports zero and circular padding, and we add "reflect" padding mode.

    Args:
        in_channels (int): Number of channels in the input feature map.
            Same as that in ``nn._ConvNd``.
        out_channels (int): Number of channels produced by the convolution.
            Same as that in ``nn._ConvNd``.
        kernel_size (int | tuple[int]): Size of the convolving kernel.
            Same as that in ``nn._ConvNd``.
        stride (int | tuple[int]): Stride of the convolution.
            Same as that in ``nn._ConvNd``.
        padding (int | tuple[int]): Zero-padding added to both sides of
            the input. Same as that in ``nn._ConvNd``.
        dilation (int | tuple[int]): Spacing between kernel elements.
            Same as that in ``nn._ConvNd``.
        groups (int): Number of blocked connections from input channels to
            output channels. Same as that in ``nn._ConvNd``.
        bias (bool | str): If specified as `auto`, it will be decided by the
            norm_cfg. Bias will be set as True if `norm_cfg` is None, otherwise
            False. Default: "auto".
        conv_cfg (dict): Config dict for convolution layer. Default: None,
            which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer. Default: None.
        act_cfg (dict): Config dict for activation layer.
            Default: dict(type='ReLU').
        inplace (bool): Whether to use inplace mode for activation.
            Default: True.
        with_spectral_norm (bool): Whether use spectral norm in conv module.
            Default: False.
        padding_mode (str): If the `padding_mode` has not been supported by
            current `Conv2d` in PyTorch, we will use our own padding layer
            instead. Currently, we support ['zeros', 'circular'] with official
            implementation and ['reflect'] with our own implementation.
            Default: 'zeros'.
        order (tuple[str]): The order of conv/norm/activation layers. It is a
            sequence of "conv", "norm" and "act". Common examples are
            ("conv", "norm", "act") and ("act", "conv", "norm").
            Default: ('conv', 'norm', 'act').
       r   autoNReLUtypeTFzerosconvnormactc                    s  t t|   |	d u st|	tsJ |
d u st|
tsJ |d u s(t|ts(J ddg}t|	| _t|
| _t|| _	|| _
|| _|| _|| _
|| _|
| _||v| _|| _t| jtret| jdksgJ t|tg dkssJ |
d u| _|d u| _|dkr| j }|| _| jrt|d}t||| _| jrdn|}t|	||||||||d	| _| jrtj| j| _| jr|d	|d
kr|}n|}t |
|\| _!}| "| j!| | jr| }|d dvr|#d
| t$|| _%| &  d S )Nr   circular   r   r
   r   r   )stridepaddingdilationgroupsbiasr   r   r
   )TanhPReLUSigmoidHSigmoidZSwishinplace)'superr   __init__
isinstancedictcopydeepcopyconv_cfgnorm_cfgact_cfgin_channelsout_channelskernel_sizer   r   with_spectral_normwith_explicit_paddingordertuplelenset	with_normwith_activationZ	with_biasr   
padding_layerr   r   nnutils
spectral_normindexr   	norm_name
add_module
setdefaultr   activateinit_weights)selfr(   r)   r*   r   r   r   r   r   r%   r&   r'   r   r+   padding_moder-   Zofficial_padding_modeZpad_cfgZconv_paddingZ
norm_channelsr   Zact_cfg_	__class__ 5/root/Awesome-Backbones/configs/common/conv_module.pyr    9   sh   




zConvModule.__init__c                 C   s   | j r	t| | j S d S )N)r8   getattr)r=   rA   rA   rB   r      s   zConvModule.normc                 C   sj   t | jds&| jr| jd dkrd}| jdd}nd}d}t| j||d	 | jr3t| jd
dd d S d S )Nr<   r
   	LeakyReLU
leaky_relunegative_slopeg{Gz?relur   )anonlinearityr	   )r   )	hasattrr   r2   r'   getr   r1   r   r   )r=   rI   rH   rA   rA   rB   r<      s   zConvModule.init_weightsc                 C   sn   | j D ]1}|dkr| jr| |}| |}q|dkr&|r&| jr&| |}q|dkr4|r4| jr4| |}q|S )Nr   r   r   )r-   r,   r3   r   r1   r   r2   r;   )r=   xr;   r   layerrA   rA   rB   forward   s   


zConvModule.forward)TT)__name__
__module____qualname____doc__r"   r    propertyr   r<   rN   
__classcell__rA   rA   r?   rB   r      s&    6R
r   )
torch.nnr4   Zconfigs.basic.build_layerr   r   r   r   Zcore.initializer   r   r#   Moduler   rA   rA   rA   rB   <module>   s
    