COCOLoss.py

# encoding:utf-8
import torch 
import torch.nn.functional as F

import torch.nn as nn
from torch.nn import Parameter

class COCOLoss(nn.Module):

    def __init__(self, num_classes, feat_dim=2048, alpha=50, use_gpu=True):
        super(COCOLoss, self).__init__()
        self.feat_dim = feat_dim
        self.num_classes = num_classes
        self.alpha = alpha
        self.centers = nn.Parameter(torch.randn(num_classes, feat_dim))
        self.use_gpu = use_gpu
        if self.use_gpu:
            self.centers = nn.Parameter(torch.randn(num_classes, feat_dim).cuda())

    def forward(self, feat):
        norms = torch.norm(feat, p=2, dim=-1, keepdim=True)
        nfeat = torch.div(feat, norms)
        snfeat = self.alpha*nfeat

        norms_c = torch.norm(self.centers, p=2, dim=-1, keepdim=True)
        ncenters = torch.div(self.centers, norms_c)

        logits = torch.matmul(snfeat, torch.transpose(ncenters, 0, 1))

        return logits