pytoch实现nn.CrossEntropyLoss和多分类的focal loss

最新推荐文章于 2025-01-16 11:15:13 发布

laizi_laizi

最新推荐文章于 2025-01-16 11:15:13 发布

阅读量1.2k

点赞数 3

分类专栏： PyTorch

本文链接：https://blog.csdn.net/laizi_laizi/article/details/117114958

版权

PyTorch 专栏收录该内容

13 篇文章

订阅专栏

本文对比了PyTorch中的nn.CrossEntropyLoss、Focal Loss的实现，并提供了自定义函数，展示了如何使用权重调整和不同reduction策略。通过代码实例演示了三种损失函数的性能和用法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

现在不太想写关于nn.CrossEntropyLoss和focal loss的解析，都还是比较容易理解的，这里直接贴上代码：

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

# for weight variable:
# https://discuss.pytorch.org/t/passing-the-weights-to-crossentropyloss-correctly/14731/10?u=simiao_lai1

def crossentropy(logits, targets, weight=None, reduction='mean'):
    """N samples, C classes
    logits: [N, C]
    targets:[N] range [0, C-1]
    weight: [C]
    """
    C = logits.size(1)
    if weight is not None:
        assert len(weight)==C, 'weight length must be equal to classes number'
        assert weight.dim() == 1, 'weight dim must be 1'

    else:
        weight = torch.ones(C)

    log_prob = F.log_softmax(logits, dim=1)
    tar_one_hot = F.one_hot(targets, num_classes=C).type(torch.float32)

    loss = weight[targets] * (-log_prob * tar_one_hot).sum(dim=1)
    if reduction == 'mean':
        loss = loss.sum() / (weight[targets].sum() + 1e-7)
    elif reduction == 'none':
        loss = loss

    return loss

def focalloss(logits, targets, gamma=0.0, weight=None, reduction='mean'):
    """N samples, C classes
    logits: [N, C]
    targets:[N] range [0, C-1]
    gamma: factor(default 0.0, that is standard cross entropy)
    weight: [C]
    """
    N, C = logits.size(0), logits.size(1)
    if weight is not None:
        assert len(weight) == C, 'weight length must be equal to classes number'
        assert weight.dim() == 1, 'weight dim must be 1'

    else:
        weight = torch.ones(C)

    prob = F.softmax(logits, dim=1)
    log_prob = F.log_softmax(logits, dim=1)
    tar_one_hot = F.one_hot(targets, num_classes=C).type(torch.float32)
    factor= (1 - prob[range(N), targets]) ** gamma

    loss = factor * weight[targets] * (-log_prob * tar_one_hot).sum(dim=1)
    if reduction == 'mean':
        loss = loss.sum() / (weight[targets].sum() + 1e-7)
    elif reduction == 'none':
        loss = loss

    return loss


N, C = 3, 6
logit = torch.randn(N, C)
target = torch.randint(0, C, (N,))
weight = torch.rand(C)

criterion = nn.CrossEntropyLoss(weight=weight, reduction='none')
s = time.time()
loss1 = criterion(logit, target)
print(time.time()-s)
s = time.time()
loss2 = crossentropy(logit, target, weight=weight, reduction='none')
print(time.time()-s)
s = time.time()
loss3 = focalloss(logit, target, gamma=0.5, weight=weight, reduction='none')
print(time.time()-s)

print(loss1)
print(loss2)
print(loss3)