| | |
| | import torch |
| | import torch.nn as nn |
| |
|
| |
|
| | class DeepLabCE(nn.Module): |
| | """ |
| | Hard pixel mining with cross entropy loss, for semantic segmentation. |
| | This is used in TensorFlow DeepLab frameworks. |
| | Paper: DeeperLab: Single-Shot Image Parser |
| | Reference: https://github.com/tensorflow/models/blob/bd488858d610e44df69da6f89277e9de8a03722c/research/deeplab/utils/train_utils.py#L33 # noqa |
| | Arguments: |
| | ignore_label: Integer, label to ignore. |
| | top_k_percent_pixels: Float, the value lies in [0.0, 1.0]. When its |
| | value < 1.0, only compute the loss for the top k percent pixels |
| | (e.g., the top 20% pixels). This is useful for hard pixel mining. |
| | weight: Tensor, a manual rescaling weight given to each class. |
| | """ |
| |
|
| | def __init__(self, ignore_label=-1, top_k_percent_pixels=1.0, weight=None): |
| | super(DeepLabCE, self).__init__() |
| | self.top_k_percent_pixels = top_k_percent_pixels |
| | self.ignore_label = ignore_label |
| | self.criterion = nn.CrossEntropyLoss( |
| | weight=weight, ignore_index=ignore_label, reduction="none" |
| | ) |
| |
|
| | def forward(self, logits, labels, weights=None): |
| | if weights is None: |
| | pixel_losses = self.criterion(logits, labels).contiguous().view(-1) |
| | else: |
| | |
| | pixel_losses = self.criterion(logits, labels) * weights |
| | pixel_losses = pixel_losses.contiguous().view(-1) |
| | if self.top_k_percent_pixels == 1.0: |
| | return pixel_losses.mean() |
| |
|
| | top_k_pixels = int(self.top_k_percent_pixels * pixel_losses.numel()) |
| | pixel_losses, _ = torch.topk(pixel_losses, top_k_pixels) |
| | return pixel_losses.mean() |
| |
|
