o

    WA-h5                     @   s<   d dl Z d dlmZ d dlm  mZ G dd dejZdS )    Nc                       s&   e Zd Zd fdd	Zdd Z  ZS )
SupConLossQ?      ?c                    s   t t|   || _|| _d S )N)superr   __init__temperatureloss_weight)selfr   r   	__class__ 5/root/Awesome-Backbones/configs/losses/supcon_loss.pyr      s   
zSupConLoss.__init__c                 C   s  |j }|jd }| dd}t||j |}|jd }tj	tj
|dddd}|}|}	tt||j| j
}
t|
tj||	 |d }||	|}|| }t|
| }|
t|jdddd  }
|
| d|d }| }||	| }|| j S )	zK 
        features: [bsz, n_views, feat_dim]
        labels: [bsz]
        r      )dim)deviceT)keepdimg-q=)r   shape
contiguousviewtorcheqTfloattocatunbinddivmatmulr   	ones_likeeyerepeatexplogsummeanr   )r	   featureslabelsr   bszmaskZcontrast_countZcontrast_featureZanchor_featureZanchor_countZanchor_dot_contrastZlogits_maskZ
exp_logitslog_probZmean_log_prob_poslossr   r   r
   forward   s*   


zSupConLoss.forward)r   r   )__name__
__module____qualname__r   r,   
__classcell__r   r   r
   r
   r      s    r   )r   torch.nnnntorch.nn.functional
functionalFModuler   r   r   r   r
   <module>   s    