| 18 | args = cfg.parse_args() |
| 19 | |
| 20 | class CombinedLoss(nn.Module): |
| 21 | def __init__(self, dice_weight=1, focal_weight=1): |
| 22 | super(CombinedLoss, self).__init__() |
| 23 | self.dice_weight = dice_weight |
| 24 | self.focal_weight = focal_weight |
| 25 | self.dice_loss = DiceLoss(to_onehot_y=True, sigmoid=True) |
| 26 | self.focal_loss = FocalLoss(to_onehot_y=True, gamma=2.0) |
| 27 | |
| 28 | def forward(self, inputs, targets): |
| 29 | dice = self.dice_loss(inputs, targets) |
| 30 | focal = self.focal_loss(inputs, targets) |
| 31 | return self.dice_weight * dice + self.focal_weight * focal |
| 32 | |
| 33 | |
| 34 | GPUdevice = torch.device('cuda', args.gpu_device) |