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   )ClsHeadc                       sJ   e Zd ZdZdeddddf fdd	Zd	d
 ZdddZd
d Z  Z	S )EfficientFormerClsHeada  EfficientFormer classifier head.

    Args:
        num_classes (int): Number of categories excluding the background
            category.
        in_channels (int): Number of channels in the input feature map.
        distillation (bool): Whether use a additional distilled head.
            Defaults to True.
        init_cfg (dict): The extra initialization configs. Defaults to
            ``dict(type='Normal', layer='Linear', std=0.01)``.
    TNormalLinearg{Gz?)typelayerstdc                    sx   t t| j|d|i| || _|| _|| _| jdkr#td| dt| j| j| _	| jr:t| j| j| _
d S d S )Ninit_cfgr   znum_classes=z must be a positive integer)superr   __init__in_channelsnum_classesdist
ValueErrornnr   head	dist_head)selfr   r
   Zdistillationr
   argskwargs	__class__ =/root/Awesome-Backbones/configs/heads/efficientformer_head.pyr      s$   


zEfficientFormerClsHead.__init__c                 C   s   t |tr	|d }|S )N)
isinstancetuple)r   xr   r   r   
pre_logits*   s   
z!EfficientFormerClsHead.pre_logitsc                 C   sb   |  |}| |}| jr|| | 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]]): 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 and cls token. The cls token will be used
                to classify and the shape of it 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   r   r   Fsoftmaxpost_process)r   r   r#   r$   	cls_scorepredr   r   r   simple_test/   s   


z"EfficientFormerClsHead.simple_testc                 K   s:   | j rtd| |}| |}| j||fi |}|S )NzPMMClassification doesn't support to train the distilled version EfficientFormer.)r   NotImplementedErrorr   r   loss)r   r   gt_labelr   r%   lossesr   r   r   
forward_trainU   s   

z$EfficientFormerClsHead.forward_train)TT)
__name__
__module____qualname____doc__dictr   r   r'   r,   
__classcell__r   r   r   r   r      s    
&r   )torch.nnr   Ztorch.nn.functional
functionalr"   cls_headr   r   r   r   r   r   <module>   s   