Gradient loss
| 108 | |
| 109 | |
| 110 | class GradL1Loss(nn.Module): |
| 111 | """Gradient loss""" |
| 112 | def __init__(self): |
| 113 | super(GradL1Loss, self).__init__() |
| 114 | self.name = 'GradL1' |
| 115 | |
| 116 | def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False): |
| 117 | input = extract_key(input, KEY_OUTPUT) |
| 118 | if input.shape[-1] != target.shape[-1] and interpolate: |
| 119 | input = nn.functional.interpolate( |
| 120 | input, target.shape[-2:], mode='bilinear', align_corners=True) |
| 121 | intr_input = input |
| 122 | else: |
| 123 | intr_input = input |
| 124 | |
| 125 | grad_gt = grad(target) |
| 126 | grad_pred = grad(input) |
| 127 | mask_g = grad_mask(mask) |
| 128 | |
| 129 | loss = nn.functional.l1_loss(grad_pred[0][mask_g], grad_gt[0][mask_g]) |
| 130 | loss = loss + \ |
| 131 | nn.functional.l1_loss(grad_pred[1][mask_g], grad_gt[1][mask_g]) |
| 132 | if not return_interpolated: |
| 133 | return loss |
| 134 | return loss, intr_input |
| 135 | |
| 136 | |
| 137 | class OrdinalRegressionLoss(object): |