Image Weighted Cross Entropy Loss
| 68 | |
| 69 | |
| 70 | class ImageBasedCrossEntropyLoss2d(nn.Module): |
| 71 | """ |
| 72 | Image Weighted Cross Entropy Loss |
| 73 | """ |
| 74 | |
| 75 | def __init__(self, classes, weight=None, ignore_index=cfg.DATASET.IGNORE_LABEL, |
| 76 | norm=False, upper_bound=1.0, fp16=False): |
| 77 | super(ImageBasedCrossEntropyLoss2d, self).__init__() |
| 78 | logx.msg("Using Per Image based weighted loss") |
| 79 | self.num_classes = classes |
| 80 | self.nll_loss = nn.NLLLoss(weight, reduction='mean', |
| 81 | ignore_index=ignore_index) |
| 82 | self.norm = norm |
| 83 | self.upper_bound = upper_bound |
| 84 | self.batch_weights = cfg.BATCH_WEIGHTING |
| 85 | self.fp16 = fp16 |
| 86 | |
| 87 | def calculate_weights(self, target): |
| 88 | """ |
| 89 | Calculate weights of classes based on the training crop |
| 90 | """ |
| 91 | bins = torch.histc(target, bins=self.num_classes, min=0.0, |
| 92 | max=self.num_classes) |
| 93 | hist_norm = bins.float() / bins.sum() |
| 94 | if self.norm: |
| 95 | hist = ((bins != 0).float() * self.upper_bound * |
| 96 | (1 / hist_norm)) + 1.0 |
| 97 | else: |
| 98 | hist = ((bins != 0).float() * self.upper_bound * |
| 99 | (1. - hist_norm)) + 1.0 |
| 100 | return hist |
| 101 | |
| 102 | def forward(self, inputs, targets, do_rmi=None): |
| 103 | |
| 104 | if self.batch_weights: |
| 105 | weights = self.calculate_weights(targets) |
| 106 | self.nll_loss.weight = weights |
| 107 | |
| 108 | loss = 0.0 |
| 109 | for i in range(0, inputs.shape[0]): |
| 110 | if not self.batch_weights: |
| 111 | weights = self.calculate_weights(targets) |
| 112 | if self.fp16: |
| 113 | weights = weights.half() |
| 114 | self.nll_loss.weight = weights |
| 115 | |
| 116 | loss += self.nll_loss(F.log_softmax(inputs[i].unsqueeze(0), dim=1), |
| 117 | targets[i].unsqueeze(0),) |
| 118 | return loss |
| 119 | |
| 120 | |
| 121 | class CrossEntropyLoss2d(nn.Module): |