o

    h$                     @   sb   d dl mZ d dlm  mZ ddlmZ ddlm	Z	 ddl
mZ ddlm
Z
 G dd	 d	e
ZdS )
    N   )
ConvModule)SELayer   )DropPath)
BaseModulec                	       sH   e Zd ZdZddddeddedddd	df	 fd
d	Zdd
 Z  ZS )InvertedResiduala  Inverted Residual Block.

    Args:
        in_channels (int): The input channels of this module.
        out_channels (int): The output channels of this module.
        mid_channels (int): The input channels of the depthwise convolution.
        kernel_size (int): The kernel size of the depthwise convolution.
            Defaults to 3.
        stride (int): The stride of the depthwise convolution. Defaults to 1.
        se_cfg (dict, optional): Config dict for se layer. Defaults to None,
            which means no se layer.
        conv_cfg (dict): Config dict for convolution layer. Defaults to None,
            which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Defaults to ``dict(type='BN')``.
        act_cfg (dict): Config dict for activation layer.
            Defaults to ``dict(type='ReLU')``.
        drop_path_rate (float): stochastic depth rate. Defaults to 0.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
            memory while slowing down the training speed. Defaults to False.
        init_cfg (dict | list[dict], optional): Initialization config dict.
       r   NBN)typeReLUg        Fc
           
         s   t t| | |dko||k| _|dv sJ || _|
dkr"t|
nt | _|d u| _	||k| _
| j	r;t|ts;J | j
rKt
||ddd|||	d| _t
|||||d ||||	d	| _| j	rftdi || _t
||ddd||d d| _d S )Nr   )r   r   r   )in_channelsout_channelskernel_sizestridepaddingconv_cfgnorm_cfgact_cfgr   )	r
   r   r   r   r   groupsr   r   r    )superr   __init__with_res_shortcutwith_cpr   nnIdentity	drop_pathwith_sewith_expand_conv
isinstancedictr   expand_convdepthwise_convr   selinear_conv)
selfr
   r   mid_channelsr   r   se_cfgr   r   r   Zdrop_path_rater   init_cfg	__class__r   ;/root/Awesome-Backbones/configs/common/inverted_residual.pyr   "   s`   


	
zInvertedResidual.__init__c                    s4    fdd} j r|jrt||}|S ||}|S )zForward function.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            torch.Tensor: The output tensor.
        c                    sP   | } j r
 |} |} jr |} |} jr&|  | S |S )N)r   r"   r#   r   r$   r%   r   r   )xoutr&   r   r,   _inner_forwarde   s   



z0InvertedResidual.forward.<locals>._inner_forward)r   
requires_gradcp
checkpoint)r&   r-   r0   r.   r   r/   r,   forward[   s   
zInvertedResidual.forward)__name__
__module____qualname____doc__r!   r   r4   
__classcell__r   r   r*   r,   r   
   s    9r   )torch.nnr   torch.utils.checkpointutilsr3   r2   conv_moduler   se_layerr   Z
basic.dropr   base_moduler   r   r   r   r   r,   <module>   s   