Spaces:
Running
Running
File size: 2,506 Bytes
123489f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import defaultdict
def dice_loss(input_mask, cls_gt):
num_objects = input_mask.shape[1]
losses = []
for i in range(num_objects):
mask = input_mask[:, i].flatten(start_dim=1)
# background not in mask, so we add one to cls_gt
gt = (cls_gt == (i + 1)).float().flatten(start_dim=1)
numerator = 2 * (mask * gt).sum(-1)
denominator = mask.sum(-1) + gt.sum(-1)
loss = 1 - (numerator + 1) / (denominator + 1)
losses.append(loss)
return torch.cat(losses).mean()
# https://stackoverflow.com/questions/63735255/how-do-i-compute-bootstrapped-cross-entropy-loss-in-pytorch
class BootstrappedCE(nn.Module):
def __init__(self, start_warm, end_warm, top_p=0.15):
super().__init__()
self.start_warm = start_warm
self.end_warm = end_warm
self.top_p = top_p
def forward(self, input, target, it):
if it < self.start_warm:
return F.cross_entropy(input, target), 1.0
raw_loss = F.cross_entropy(input, target, reduction="none").view(-1)
num_pixels = raw_loss.numel()
if it > self.end_warm:
this_p = self.top_p
else:
this_p = self.top_p + (1 - self.top_p) * (
(self.end_warm - it) / (self.end_warm - self.start_warm)
)
loss, _ = torch.topk(raw_loss, int(num_pixels * this_p), sorted=False)
return loss.mean(), this_p
class LossComputer:
def __init__(self, config):
super().__init__()
self.config = config
self.bce = BootstrappedCE(config["start_warm"], config["end_warm"])
def compute(self, data, num_objects, it):
losses = defaultdict(int)
b, t = data["rgb"].shape[:2]
losses["total_loss"] = 0
for ti in range(1, t):
for bi in range(b):
loss, p = self.bce(
data[f"logits_{ti}"][bi : bi + 1, : num_objects[bi] + 1],
data["cls_gt"][bi : bi + 1, ti, 0],
it,
)
losses["p"] += p / b / (t - 1)
losses[f"ce_loss_{ti}"] += loss / b
losses["total_loss"] += losses["ce_loss_%d" % ti]
losses[f"dice_loss_{ti}"] = dice_loss(
data[f"masks_{ti}"], data["cls_gt"][:, ti, 0]
)
losses["total_loss"] += losses[f"dice_loss_{ti}"]
return losses
|