text: # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import paddle
  import paddle.nn as nn
  
  from .builder import HEADS
  
  
  @HEADS.register()
  class ContrastiveHead(nn.Layer):
      """Head for contrastive learning.
  
      Args:
          temperature (float): The temperature hyper-parameter that
              controls the concentration level of the distribution.
              Default: 0.1.
      """
  
      def __init__(self, temperature=0.1, return_accuracy=True):
          super(ContrastiveHead, self).__init__()
          self.criterion = nn.CrossEntropyLoss()
          self.temperature = temperature
          self.return_accuracy = return_accuracy
  
      def forward(self, pos, neg):
          """Forward head.
  
          Args:
              pos (Tensor): Nx1 positive similarity.
              neg (Tensor): Nxk negative similarity.
  
          Returns:
              dict[str, Tensor]: A dictionary of loss components.
          """
          N = pos.shape[0]
          logits = paddle.concat((pos, neg), axis=1)
          logits /= self.temperature
          labels = paddle.zeros((N, ), dtype='int64')
          outputs = dict()
          outputs['loss'] = self.criterion(logits, labels)
          acc1, acc5 = accuracy(logits, labels, topk=(1, 5))
          outputs['acc1'] = acc1
          outputs['acc5'] = acc5
          return outputs
  
  
  def accuracy(output, target, topk=(1, )):
      """Computes the accuracy over the k top predictions for the specified values of k"""
      with paddle.no_grad():
          maxk = max(topk)
          batch_size = target.shape[0]
  
          _, pred = output.topk(maxk, 1, True, True)
          pred = pred.t()
          correct = paddle.cast(pred == target.reshape([1, -1]).expand_as(pred),
                                'float32')
  
          res = []
          for k in topk:
              correct_k = correct[:k].reshape([-1]).sum(0, keepdim=True)
              res.append(correct_k * 100.0 / batch_size)
          return res
  