Source code for catalyst.contrib.criterion.ce
import torch
import torch.nn as nn
import torch.nn.functional as F
[docs]class NaiveCrossEntropyLoss(nn.Module):
def __init__(self, size_average=True):
super().__init__()
self.size_average = size_average
[docs] def forward(self, input, target):
assert input.size() == target.size()
input = F.log_softmax(input)
loss = -torch.sum(input * target)
loss = loss / input.size()[0] if self.size_average else loss
return loss
[docs]class MaskCrossEntropyLoss(torch.nn.CrossEntropyLoss):
def __init__(
self,
*args,
target_name: str = "targets",
mask_name: str = "mask",
**kwargs
):
super().__init__(*args, **kwargs)
self.target_name = target_name
self.mask_name = mask_name
self.reduction = "none"
[docs] def forward(self, input, target_mask):
target = target_mask[self.target_name]
mask = target_mask[self.mask_name]
loss = super().forward(input, target)
loss = torch.mean(loss[mask == 1])
return loss
__all__ = ["MaskCrossEntropyLoss", "NaiveCrossEntropyLoss"]