o

    h                     @   sb   d dl Z d dlZd dlmZ d dlm  mZ ddlT d dlm	Z	 ddl
mZ G dd deZdS )    N   )*)Accuracy)
BaseModulec                       sT   e Zd ZdZeddddddf fdd		Zd
d Zdd
 Zdd ZdddZ	  Z
S )ClsHeadab  classification head.

    Args:
        loss (dict): Config of classification loss.
        topk (int | tuple): Top-k accuracy.
        cal_acc (bool): Whether to calculate accuracy during training.
            If you use Mixup/CutMix or something like that during training,
            it is not reasonable to calculate accuracy. Defaults to False.
    CrossEntropyLossg      ?)typeZloss_weight)   FNc                    s   t t| j|d t|tsJ t|ttfsJ t|tr!|f}|D ]
}|dks-J dq#|| _t|	ddi || _
t| jd| _|| _
d S )N)init_cfgr   zTop-k should be larger than 0r   )topk )superr   __init__
isinstancedictinttupler   evalpopcompute_lossr   compute_accuracycal_acc)selflossr   r   r
   Z_topk	__class__r   1/root/Awesome-Backbones/configs/heads/cls_head.pyr      s   

zClsHead.__init__c                 K   s   t |}t }| j||fd|i|}| jrHt & | ||}t |t | jks,J dd t| j|D |d< W d    n1 sCw   Y  ||d< |S )NZ
avg_factorc                 S   s   i | ]
\}}d | |qS )ztop-r   ).0kar   r   r   
<dictcomp>7   s    
z ClsHead.loss.<locals>.<dictcomp>accuracyr   )	lenr   r   r   torchno_gradr   r   zip)r   	cls_scoregt_labelkwargsnum_sampleslossesr   accr   r   r   r   ,   s&   

zClsHead.lossc                 K   s*   t |tr	|d }| j||fi |}|S )N)r   r   r   )r   r&   r'   r(   r*   r   r   r   
forward_train?   s   
zClsHead.forward_trainc                 C   s    t |tr	|d }td |S )Nr,   zpThe input of ClsHead should be already logits. Please modify the backbone if you want to get pre-logits feature.)r   r   warningswarn)r   xr   r   r   
pre_logitsE   s   
zClsHead.pre_logitsTc                 C   sH   t |tr	|d }|r|durtj|ddnd}n|}|r"| |S |S )a  Inference without augmentation.

        Args:
            cls_score (tuple[Tensor]): The input classification score logits.
                Multi-stage inputs are acceptable but only the last stage will
                be used to classify. The shape of every item should be
                ``(num_samples, num_classes)``.
            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)``.
        r,   Nr	   )dim)r   r   Fsoftmaxpost_process)r   r&   r4   r5   predr   r   r   simple_testO   s   

zClsHead.simple_test)TF)__name__
__module____qualname____doc__r   r   r   r-   r1   r7   
__classcell__r   r   r   r   r      s    

r   )
r.   r#   torch.nnnntorch.nn.functional
functionalr3   r*   core.evaluationsr   Zcommon.base_moduler   r   r   r   r   r   <module>   s   