def build_criterion(self):
weight_tensor = torch.FloatTensor(self.class_weights)
if not self.fp16:
return torch.nn.CrossEntropyLoss(weight=weight_tensor,
reduction='none')
else:
# FP16 safe cross entropy (softmax done in FP32)
return FP16SafeCrossEntropy(weight=weight_tensor, reduction='none')
def build_criterion(self):
"""
Construct and return the loss function.
By default torch.nn.CrossEntropyLoss.
"""
if self.fp16:
return FP16SafeCrossEntropy(reduction='none')
else:
return torch.nn.CrossEntropyLoss(reduction='none')
def build_criterion(self):
"""
Construct and return the loss function.
By default torch.nn.CrossEntropyLoss.
If overridden, this model should produce a sum that can be used for a per-token loss.
"""
if not self.fp16:
return torch.nn.CrossEntropyLoss(
ignore_index=self.NULL_IDX, reduction='none'
)
else:
# FP16 safe cross entropy (softmax done in FP32)
return FP16SafeCrossEntropy(ignore_index=self.NULL_IDX, reduction='none')