from .base_weight_regularizer import BaseWeightRegularizer
import torch

# modified from http://kaizhao.net/regularface
class RegularFaceRegularizer(BaseWeightRegularizer):
  
    def compute_loss(self, weights):
        num_classes = weights.size(0)
        cos = torch.mm(weights, weights.t())
        if not self.normalize_weights:
            norms = torch.norm(weights, p=2, dim=1, keepdim=True)
            cos = cos / (norms*norms.t())

        cos1 = cos.clone()
        with torch.no_grad():
            row_nums = torch.arange(num_classes).long().to(weights.device)
            cos1[row_nums, row_nums] = -float('inf')
            _, indices = torch.max(cos1, dim=1)
            mask = torch.zeros((num_classes, num_classes)).to(weights.device)
            mask[row_nums, indices] = 1
        
        return torch.sum(cos*mask) / num_classes