|
|
import torch |
|
|
import torch.nn as nn |
|
|
from .layers import Decoder |
|
|
import torch.nn.functional as F |
|
|
from bert.modeling_bert import BertModel |
|
|
|
|
|
def dice_loss(inputs, targets): |
|
|
""" |
|
|
Compute the DICE loss, similar to generalized IOU for masks |
|
|
Args: |
|
|
inputs: A float tensor of arbitrary shape. |
|
|
The predictions for each example. |
|
|
targets: A float tensor with the same shape as inputs. Stores the binary |
|
|
classification label for each element in inputs |
|
|
(0 for the negative class and 1 for the positive class). |
|
|
""" |
|
|
|
|
|
inputs = inputs.sigmoid() |
|
|
inputs = inputs.flatten(1) |
|
|
targets = targets.flatten(1) |
|
|
numerator = 2 * (inputs * targets).sum(1) |
|
|
denominator = inputs.sum(-1) + targets.sum(-1) |
|
|
loss = 1 - (numerator + 1) / (denominator + 1) |
|
|
return loss.mean() |
|
|
|
|
|
def sigmoid_focal_loss(inputs, targets, alpha: float = 0.25, gamma: float = 2): |
|
|
""" |
|
|
Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. |
|
|
Args: |
|
|
inputs: A float tensor of arbitrary shape. |
|
|
The predictions for each example. |
|
|
targets: A float tensor with the same shape as inputs. Stores the binary |
|
|
classification label for each element in inputs |
|
|
(0 for the negative class and 1 for the positive class). |
|
|
alpha: (optional) Weighting factor in range (0,1) to balance |
|
|
positive vs negative examples. Default = -1 (no weighting). |
|
|
gamma: Exponent of the modulating factor (1 - p_t) to |
|
|
balance easy vs hard examples. |
|
|
Returns: |
|
|
Loss tensor |
|
|
""" |
|
|
|
|
|
prob = inputs.sigmoid() |
|
|
ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") |
|
|
p_t = prob * targets + (1 - prob) * (1 - targets) |
|
|
loss = ce_loss * ((1 - p_t) ** gamma) |
|
|
|
|
|
if alpha >= 0: |
|
|
alpha_t = alpha * targets + (1 - alpha) * (1 - targets) |
|
|
loss = alpha_t * loss |
|
|
return loss.mean() |
|
|
|
|
|
|
|
|
class CGFormer(nn.Module): |
|
|
def __init__(self, backbone, args): |
|
|
super(CGFormer, self).__init__() |
|
|
self.backbone = backbone |
|
|
self.decoder = Decoder(args) |
|
|
self.text_encoder = BertModel.from_pretrained(args.bert) |
|
|
self.text_encoder.pooler = None |
|
|
|
|
|
def forward(self, x, text, l_mask, mask=None): |
|
|
input_shape = x.shape[-2:] |
|
|
l_feats = self.text_encoder(text, attention_mask=l_mask)[0] |
|
|
l_feats = l_feats.permute(0, 2, 1) |
|
|
l_mask = l_mask.unsqueeze(dim=-1) |
|
|
|
|
|
|
|
|
features = self.backbone(x, l_feats, l_mask) |
|
|
x_c1, x_c2, x_c3, x_c4 = features |
|
|
pred, maps = self.decoder([x_c4, x_c3, x_c2, x_c1], l_feats, l_mask) |
|
|
pred = F.interpolate(pred, input_shape, mode='bilinear', align_corners=True) |
|
|
|
|
|
|
|
|
if self.training: |
|
|
loss = 0. |
|
|
mask = mask.unsqueeze(1).float() |
|
|
for m, lam in zip(maps, [0.001,0.01,0.1]): |
|
|
m = m[:,1].unsqueeze(1) |
|
|
if m.shape[-2:] != mask.shape[-2:]: |
|
|
mask_ = F.interpolate(mask, m.shape[-2:], mode='nearest').detach() |
|
|
loss += dice_loss(m, mask_) * lam |
|
|
loss += dice_loss(pred, mask) + sigmoid_focal_loss(pred, mask, alpha=-1, gamma=0) |
|
|
return pred.detach(), mask, loss |
|
|
else: |
|
|
return pred.detach(), maps |
