focal loss与batch-balanced contrastive loss 相结合pytorch代码

下面是结合 Focal Loss 和 Batch-Balanced Contrastive Loss 的 PyTorch 代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
        
    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss

class BatchBalancedContrastiveLoss(nn.Module):
    def __init__(self, margin=0.5, alpha=0.5, beta=1, gamma=2, reduction='mean'):
        super(BatchBalancedContrastiveLoss, self).__init__()
        self.margin = margin
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.reduction = reduction
        
    def forward(self, inputs, targets):
        n = inputs.size(0)
        sim_mat = torch.matmul(inputs, inputs.t())
        
        targets = targets.view(n,1)
        mask = targets.expand(n,n).eq(targets.expand(n,n).t())
        
        pos_mask = mask.triu(diagonal=1)
        neg_mask = (mask-triu(diagonal=1)).bool()
        
        pos_pair = sim_mat[pos_mask]
        neg_pair = sim_mat[neg_mask]
        
        num_pos_pair = pos_mask.sum()
        num_neg_pair = neg_mask.sum()
        
        alpha = self.alpha
        beta = self.beta
        
        if num_pos_pair > 0:
            alpha = (num_neg_pair / num_pos_pair) * self.alpha
        
        if num_neg_pair > 0:
            beta = (num_pos_pair / num_neg_pair) * self.beta
        
        pos_loss = F.relu(pos_pair - self.margin)
        neg_loss = F.relu(self.margin - neg_pair)
        
        if self.gamma > 0:
            pos_loss = torch.pow(pos_loss, self.gamma)
            neg_loss = torch.pow(neg_loss, self.gamma)
        
        pos_loss = alpha * pos_loss
        neg_loss = beta * neg_loss
        
        bbcon_loss = torch.cat([pos_loss, neg_loss], dim=0)
        
        if self.reduction == 'mean':
            return bbcon_loss.mean()
        elif self.reduction == 'sum':
            return bbcon_loss.sum()
        else:
            return bbcon_loss

class FocalBatchBalancedContrastiveLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2, margin=0.5, beta=1, reduction='mean'):
        super(FocalBatchBalancedContrastiveLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.margin = margin
        self.beta = beta
        self.reduction = reduction
        
        self.focal_loss = FocalLoss(alpha=self.alpha, gamma=self.gamma, reduction='none')
        self.bbcon_loss = BatchBalancedContrastiveLoss(margin=self.margin, alpha=1, beta=self.beta, reduction='none')
        
    def forward(self, inputs, targets):
        n = inputs.size(0)
        ce_loss = self.focal_loss(inputs, targets)
        bbcon_loss = self.bbcon_loss(inputs, targets)
        
        if self.reduction == 'mean':
            return (ce_loss + bbcon_loss).mean()
        elif self.reduction == 'sum':
            return (ce_loss + bbcon_loss).sum()
        else:
            return ce_loss + bbcon_loss

使用方法：

loss_fn = FocalBatchBalancedContrastiveLoss(alpha=1, gamma=2, margin=0.5, beta=1, reduction='mean')
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
        loss.backward()
        optimizer.step()

注意，这里的 `inputs` 是模型的输出，而 `targets` 是标签。在训练过程中，你需要将 `inputs` 和 `targets` 输入到模型中，得到输出 `outputs`，然后计算损失并进行反向传播和优化。