o

    h                     @   s>   d dl mZ d dlm  mZ ddlmZ G dd deZdS )    N   )VisionTransformerClsHeadc                       s6   e Zd Z fddZdd ZdddZd	d
 Z  ZS )DeiTClsHeadc                    sT   t t| j|i | | jd u rt| j| j}nt| j| j}| j	d| d S )N	head_dist)
superr   __init__
hidden_dimnnLinearin_channelsnum_classeslayers
add_module)selfargskwargsr   	__class__ 2/root/Awesome-Backbones/configs/heads/deit_head.pyr   
   s
   
zDeiTClsHead.__init__c                 C   s^   t |tr	|d }|\}}}| jd u r||fS | j| j|}| j| j|}||fS )N)
isinstancetupler   r
   act
pre_logits)r   x_	cls_token
dist_tokenr   r   r   r      s   


zDeiTClsHead.pre_logitsTFc                 C   s`   |  |\}}| j|| j| d }|r%|dur"tj|ddnd}n|}|r.| |S |S )a$  Inference without augmentation.

        Args:
            x (tuple[tuple[tensor, tensor, tensor]]): The input features.
                Multi-stage inputs are acceptable but only the last stage will
                be used to classify. Every item should be a tuple which
                includes patch token, cls token and dist token. The cls token
                and dist token will be used to classify and the shape of them
                should be ``(num_samples, in_channels)``.
            softmax (bool): Whether to softmax the classification score.
            post_process (bool): Whether to do post processing the
                inference results. It will convert the output to a list.

        Returns:
            Tensor | list: The inference results.

                - If no post processing, the output is a tensor with shape
                  ``(num_samples, num_classes)``.
                - If post processing, the output is a multi-dimentional list of
                  float and the dimensions are ``(num_samples, num_classes)``.
           Nr   )dim)r   r
   headr   Fsoftmaxpost_process)r   r   r#   r$   r   r   	cls_scorepredr   r   r   simple_test   s"   


zDeiTClsHead.simple_testc                 C   s:   |  |\}}| j|| j| d }| ||}|S )Nr   )r   r
   r!   r   loss)r   r   gt_labelr   r   r%   lossesr   r   r   
forward_trainC   s   

zDeiTClsHead.forward_train)TF)__name__
__module____qualname__r   r   r'   r+   
__classcell__r   r   r   r   r      s
    
%r   )torch.nnr	   Ztorch.nn.functional
functionalr"   vision_transformer_headr   r   r   r   r   r   <module>   s   