o

    h	                     @   sd   d dl 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	d
 Z
G dd deZdS )
    N   )
BaseModule   )ClsHead)build_norm_layerc                       s>   e Zd Zeddf fdd	Ze dd Zdd Z  Z	S )	BatchNormLinearBN1d)typec                    s,   t t|   t||| _t||| _d S N)superr   __init__r   bnnnLinearlinear)selfin_channelsout_channelsnorm_cfg	__class__ 3/root/Awesome-Backbones/configs/heads/levit_head.pyr      s   zBatchNormLinear.__init__c                 C   s   | j j| j j| j j d  }| j j| j j| j j | j j| j j d   }| jj|d d d f  }| jj|d d d f  d| jj }| jjj	| | jjj	| | jS )Ng      ?)
r
   weightrunning_varepsbiasrunning_meanr   viewdatacopy_)r   wbr   r   r   fuse   s   &zBatchNormLinear.fusec                 C   s   |  |}| |}|S r
   )r
   r   r   xr   r   r   forward   s   

zBatchNormLinear.forward)
__name__
__module____qualname__dictr   torchno_gradr$   r'   
__classcell__r   r   r   r   r   
   s
    
r   c                 C   s:   |   D ]\}}t|drt| ||  qt| qd S )Nr$   )named_childrenhasattrsetattrr$   fuse_parameters)module
child_namechildr   r   r   r2   #   s
   

r2   c                       s6   e Zd Z				d fdd	Zdd Zd	d
 Z  ZS )LeViTClsHead  TNFc                    s\   t t| jdi | || _|| _|| _t||| _|r"t||| _| jr,| 	|  d S d S )Nr   )
r   r6   r   num_classesdistillationdeployr   head	head_distswitch_to_deploy)r   r8   r9   r   r:   kwargsr   r   r   r   -   s   zLeViTClsHead.__init__c                 C   s   | j rd S t|  d| _ d S )NT)r:   r2   )r   r   r   r   r=   >   s   
zLeViTClsHead.switch_to_deployc                 C   sX   |  |}| jr%| || |f}| js!|d |d  d }|S td| |}|S )Nr   r   r   z9MMPretrain doesn't support training in distillation mode.)
pre_logitsr9   r;   r<   trainingNotImplementedErrorr%   r   r   r   r'   D   s   

zLeViTClsHead.forward)r7   TNF)r(   r)   r*   r   r=   r'   r.   r   r   r   r   r6   +   s    r6   )r,   torch.nnr   Zcommon.base_moduler   linear_headr   configs.basic.build_layerr   r   r2   r6   r   r   r   r   <module>   s   