| 52 | |
| 53 | |
| 54 | class CGFormer(nn.Module): |
| 55 | def __init__(self, backbone, args): |
| 56 | super(CGFormer, self).__init__() |
| 57 | self.backbone = backbone |
| 58 | self.decoder = Decoder(args) |
| 59 | self.text_encoder = BertModel.from_pretrained(args.bert) |
| 60 | self.text_encoder.pooler = None |
| 61 | |
| 62 | def forward(self, x, text, l_mask, mask=None): |
| 63 | input_shape = x.shape[-2:] |
| 64 | l_feats = self.text_encoder(text, attention_mask=l_mask)[0] # (6, 10, 768) |
| 65 | l_feats = l_feats.permute(0, 2, 1) # (B, 768, N_l) to make Conv1d happy |
| 66 | l_mask = l_mask.unsqueeze(dim=-1) # (batch, N_l, 1) |
| 67 | ########################## |
| 68 | features = self.backbone(x, l_feats, l_mask) |
| 69 | x_c1, x_c2, x_c3, x_c4 = features |
| 70 | pred, maps = self.decoder([x_c4, x_c3, x_c2, x_c1], l_feats, l_mask) |
| 71 | pred = F.interpolate(pred, input_shape, mode='bilinear', align_corners=True) |
| 72 | # loss |
| 73 | if self.training: |
| 74 | loss = 0. |
| 75 | mask = mask.unsqueeze(1).float() |
| 76 | for m, lam in zip(maps, [0.001,0.01,0.1]): |
| 77 | m = m[:,1].unsqueeze(1) |
| 78 | if m.shape[-2:] != mask.shape[-2:]: |
| 79 | mask_ = F.interpolate(mask, m.shape[-2:], mode='nearest').detach() |
| 80 | loss += dice_loss(m, mask_) * lam |
| 81 | loss += dice_loss(pred, mask) + sigmoid_focal_loss(pred, mask, alpha=-1, gamma=0) |
| 82 | return pred.detach(), mask, loss |
| 83 | else: |
| 84 | return pred.detach(), maps |