mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
24 lines
862 B
Python
24 lines
862 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class DenseBinaryCrossEntropy(nn.Module):
|
|
""" BCE using one-hot from dense targets w/ label smoothing
|
|
NOTE for experiments comparing CE to BCE /w label smoothing, may remove
|
|
"""
|
|
def __init__(self, smoothing=0.1):
|
|
super(DenseBinaryCrossEntropy, self).__init__()
|
|
assert 0. <= smoothing < 1.0
|
|
self.smoothing = smoothing
|
|
self.bce = nn.BCEWithLogitsLoss()
|
|
|
|
def forward(self, x, target):
|
|
num_classes = x.shape[-1]
|
|
off_value = self.smoothing / num_classes
|
|
on_value = 1. - self.smoothing + off_value
|
|
target = target.long().view(-1, 1)
|
|
target = torch.full(
|
|
(target.size()[0], num_classes), off_value, device=x.device, dtype=x.dtype).scatter_(1, target, on_value)
|
|
return self.bce(x, target)
|