| |
| 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() |
|
|
