focal mse loss 代码

实现示例：

import torch.nn as nn

class FocalMSELoss(nn.Module):
    def __init__(self, gamma=2, reduction='mean'):
        super(FocalMSELoss, self).__init__()
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        mse_loss = nn.MSELoss(reduction=self.reduction)(inputs, targets)
        pt = torch.exp(-mse_loss)
        focal_loss = (1 - pt)**self.gamma * mse_loss
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss

您如有需要，可以根据实际情况进行修改。