o

    0-,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	T G dd deZ
G dd	 d	e
ZdS )
    N   )ClsHead)ContrastiveHead)*c                       sH   e Zd ZdZeddddf fdd	Zdd	 Zddd
Zdd Z  Z	S )
LinearClsHeada^  Linear classifier head.

    Args:
        num_classes (int): Number of categories excluding the background
            category.
        in_channels (int): Number of channels in the input feature map.
        init_cfg (dict | optional): The extra init config of layers.
            Defaults to use dict(type='Normal', layer='Linear', std=0.01).
    NormalLinearg{Gz?)typelayerstdc                    sV   t t| j|d|i| || _|| _| jdkr td| dt| j| j| _d S )Ninit_cfgr   znum_classes=z must be a positive integer)	superr   __init__in_channelsnum_classes
ValueErrornnr   fc)selfr   r   r   argskwargs	__class__ 4/root/Awesome-Backbones/configs/heads/linear_head.pyr      s   

zLinearClsHead.__init__c                 C   s   t |tr	|d }|S )N)
isinstancetuple)r   xr   r   r   
pre_logits&   s   
zLinearClsHead.pre_logitsTFc                 C   sJ   |  |}| |}|r|durtj|ddnd}n|}|r#| |S |S )af  Inference without augmentation.

        Args:
            x (tuple[Tensor]): The input features.
                Multi-stage inputs are acceptable but only the last stage will
                be used to classify. The shape of every item 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   Fsoftmaxpost_process)r   r   r"   r#   	cls_scorepredr   r   r   simple_test+   s   


zLinearClsHead.simple_testc                 K   s,   |  |}| |}| j||fi |}|S N)r   r   loss)r   r   gt_labelr   r$   lossesr   r   r   
forward_trainM   s   

zLinearClsHead.forward_train)TF)
__name__
__module____qualname____doc__dictr   r   r&   r+   
__classcell__r   r   r   r   r   
   s    

"r   c                       s@   e Zd Zddedddedddf fdd		Zd
ddZ  ZS )MultiLossLinearClsHead   gffffff?CrossEntropyLossg      ?)r	   Zloss_weightProjectedContrastiveLossg      ?c                    sL   t  jd||d| t||d| _|| _|| _t | _t	|d| _
d S )N)r   r   )in_dimproj_dim)temperaturer   )r
   r   r   contrastive_headZloss_cls_cfgZloss_contrastive_cfgr   r4   ce_lossr5   contrastive_loss)r   r   r   r7   r8   loss_clsZloss_contrastiver   r   r   r   r   T   s   
zMultiLossLinearClsHead.__init__Nc           	      K   sd   |  |}| |}t }| |||d< |d ur0|  |}| |}| |}| |||d< |S )Nr<   
loss_contrast)r   r   r0   r:   r9   r;   )	r   r   r)   x2r   r$   r*   z1z2r   r   r   r+   e   s   




z$MultiLossLinearClsHead.forward_trainr'   )r,   r-   r.   r0   r   r+   r1   r   r   r   r   r2   S   s    

r2   )torch.nnr   Ztorch.nn.functional
functionalr!   cls_headr   r9   r   configs.lossesr   r2   r   r   r   r   <module>   s   I