Cross entropy loss
| 181 | |
| 182 | |
| 183 | class DiscreteNLLLoss(nn.Module): |
| 184 | """Cross entropy loss""" |
| 185 | def __init__(self, min_depth=1e-3, max_depth=10, depth_bins=64): |
| 186 | super(DiscreteNLLLoss, self).__init__() |
| 187 | self.name = 'CrossEntropy' |
| 188 | self.ignore_index = -(depth_bins + 1) |
| 189 | # self._loss_func = nn.NLLLoss(ignore_index=self.ignore_index) |
| 190 | self._loss_func = nn.CrossEntropyLoss(ignore_index=self.ignore_index) |
| 191 | self.min_depth = min_depth |
| 192 | self.max_depth = max_depth |
| 193 | self.depth_bins = depth_bins |
| 194 | self.alpha = 1 |
| 195 | self.zeta = 1 - min_depth |
| 196 | self.beta = max_depth + self.zeta |
| 197 | |
| 198 | def quantize_depth(self, depth): |
| 199 | # depth : N1HW |
| 200 | # output : NCHW |
| 201 | |
| 202 | # Quantize depth log-uniformly on [1, self.beta] into self.depth_bins bins |
| 203 | depth = torch.log(depth / self.alpha) / np.log(self.beta / self.alpha) |
| 204 | depth = depth * (self.depth_bins - 1) |
| 205 | depth = torch.round(depth) |
| 206 | depth = depth.long() |
| 207 | return depth |
| 208 | |
| 209 | |
| 210 | |
| 211 | def _dequantize_depth(self, depth): |
| 212 | """ |
| 213 | Inverse of quantization |
| 214 | depth : NCHW -> N1HW |
| 215 | """ |
| 216 | # Get the center of the bin |
| 217 | |
| 218 | |
| 219 | |
| 220 | |
| 221 | def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False): |
| 222 | input = extract_key(input, KEY_OUTPUT) |
| 223 | # assert torch.all(input <= 0), "Input should be negative" |
| 224 | |
| 225 | if input.shape[-1] != target.shape[-1] and interpolate: |
| 226 | input = nn.functional.interpolate( |
| 227 | input, target.shape[-2:], mode='bilinear', align_corners=True) |
| 228 | intr_input = input |
| 229 | else: |
| 230 | intr_input = input |
| 231 | |
| 232 | # assert torch.all(input)<=1) |
| 233 | if target.ndim == 3: |
| 234 | target = target.unsqueeze(1) |
| 235 | |
| 236 | target = self.quantize_depth(target) |
| 237 | if mask is not None: |
| 238 | if mask.ndim == 3: |
| 239 | mask = mask.unsqueeze(1) |
| 240 |