| 99 | |
| 100 | class ClassificationLoss(torch.nn.Module): |
| 101 | def __init__(self, label_size, class_weight=None, |
| 102 | loss_type=LossType.SOFTMAX_CROSS_ENTROPY): |
| 103 | super(ClassificationLoss, self).__init__() |
| 104 | self.label_size = label_size |
| 105 | self.loss_type = loss_type |
| 106 | if loss_type == LossType.SOFTMAX_CROSS_ENTROPY: |
| 107 | self.criterion = torch.nn.CrossEntropyLoss(class_weight) |
| 108 | elif loss_type == LossType.SOFTMAX_FOCAL_CROSS_ENTROPY: |
| 109 | self.criterion = FocalLoss(label_size, ActivationType.SOFTMAX) |
| 110 | elif loss_type == LossType.SIGMOID_FOCAL_CROSS_ENTROPY: |
| 111 | self.criterion = FocalLoss(label_size, ActivationType.SIGMOID) |
| 112 | elif loss_type == LossType.BCE_WITH_LOGITS: |
| 113 | self.criterion = torch.nn.BCEWithLogitsLoss() |
| 114 | else: |
| 115 | raise TypeError( |
| 116 | "Unsupported loss type: %s. Supported loss type is: %s" % ( |
| 117 | loss_type, LossType.str())) |
| 118 | |
| 119 | def forward(self, logits, target, |
| 120 | use_hierar=False, |