| 123456789101112131415161718192021222324252627282930313233343536373839 |
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from .cross_entropy import LabelSmoothingCrossEntropy
- class JsdCrossEntropy(nn.Module):
- """ Jensen-Shannon Divergence + Cross-Entropy Loss
- Based on impl here: https://github.com/google-research/augmix/blob/master/imagenet.py
- From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
- https://arxiv.org/abs/1912.02781
- Hacked together by / Copyright 2020 Ross Wightman
- """
- def __init__(self, num_splits=3, alpha=12, smoothing=0.1):
- super().__init__()
- self.num_splits = num_splits
- self.alpha = alpha
- if smoothing is not None and smoothing > 0:
- self.cross_entropy_loss = LabelSmoothingCrossEntropy(smoothing)
- else:
- self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
- def __call__(self, output, target):
- split_size = output.shape[0] // self.num_splits
- assert split_size * self.num_splits == output.shape[0]
- logits_split = torch.split(output, split_size)
- # Cross-entropy is only computed on clean images
- loss = self.cross_entropy_loss(logits_split[0], target[:split_size])
- probs = [F.softmax(logits, dim=1) for logits in logits_split]
- # Clamp mixture distribution to avoid exploding KL divergence
- logp_mixture = torch.clamp(torch.stack(probs).mean(axis=0), 1e-7, 1).log()
- loss += self.alpha * sum([F.kl_div(
- logp_mixture, p_split, reduction='batchmean') for p_split in probs]) / len(probs)
- return loss
|
