revert dice loss

models/losses.py

@@ -165,8 +165,12 @@ def tversky_loss(true, logits, alpha, beta, eps=1e-7):
     return (1 - tversky_loss)
 
 
-def ce_dice(true, pred, log=False, w1=1, w2=1):
-    pass
+def ce_dice(true, pred, weights=torch.tensor([0.5, 2])):
+    if weights is not None:
+        weights = torch.tensor(weights).to(pred.device)
+
+    return ce_loss(true, pred, weights) + \
+           dice_loss(true, pred)
 
 
 def ce_jaccard(true, pred, weights=torch.tensor([0.5, 2])):

models/networks.py

@@ -115,8 +115,6 @@ def define_loss(opt):
     if opt.dataset_mode == 'classification':
         loss = torch.nn.CrossEntropyLoss()
     elif opt.dataset_mode == 'segmentation':
-        # loss = torch.nn.CrossEntropyLoss(ignore_index=-1, weight=torch.tensor([0.5, 2]))
-
         device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
         weights = torch.FloatTensor(opt.loss_weights).to(device)
 

train_pl.py

@@ -27,7 +27,7 @@ def __init__(self, opt):
         if opt.from_pretrained is not None:
             print('Loaded pretrained weights:', opt.from_pretrained)
             self.model.load_weights(opt.from_pretrained)
-        self.criterion = ce_jaccard
+        self.criterion = self.model.criterion
         if self.training:
             self.train_metrics = torch.nn.ModuleList([
                 torchmetrics.Accuracy(num_classes=opt.nclasses, average='macro'),
@@ -40,46 +40,34 @@ def __init__(self, opt):
                 torchmetrics.F1(num_classes=opt.nclasses, average='macro')
             ])
 
-    def training_step(self, batch, idx):
+    def step(self, batch, is_train=True):
         self.model.set_input(batch)
         out = self.model.forward()
         true, pred = postprocess(self.model.labels, out)
-        loss = self.criterion(true, pred, self.opt.class_weights)
+        loss = self.criterion(true, pred)
 
-        pred_class = out.data.max(1)[1]
-        not_padding = self.model.labels != -1
-        label_class = self.model.labels[not_padding]
-        pred_class = pred_class[not_padding]
+        true = true.view(-1)
+        pred = pred.argmax(1).view(-1)
 
+        prefix = '' if is_train else 'val_'
         for m in self.train_metrics:
-            val = m(pred_class, label_class)
+            val = m(pred, true)
             metric_name = str(m).split('(')[0]
-            self.log(metric_name.lower(), val, logger=True, prog_bar=True, on_epoch=True)
-        self.log('loss', loss, on_epoch=True)
+            self.log(prefix + metric_name.lower(), val, logger=True, prog_bar=True, on_epoch=True)
+        self.log(prefix + 'loss', loss, on_epoch=True)
         return loss
 
-    def validation_step(self, batch, idx):
-        self.model.set_input(batch)
-        out = self.model.forward()
-        true, pred = postprocess(self.model.labels, out)
-        loss = self.criterion(true, pred, self.opt.class_weights)
+    def training_step(self, batch, idx):
 
-        pred_class = out.data.max(1)[1]
-        not_padding = self.model.labels != -1
-        label_class = self.model.labels[not_padding]
-        pred_class = pred_class[not_padding]
+        return self.step(batch, is_train=True)
 
-        for m in self.val_metrics:
-            val = m(pred_class, label_class)
-            metric_name = str(m).split('(')[0]
-            self.log('val_' + metric_name.lower(), val, logger=True, prog_bar=True, on_epoch=True)
-        self.log('val_loss', loss, on_epoch=True)
-        return loss
+    def validation_step(self, batch, idx):
+        return self.step(batch, is_train=False)
 
     def forward(self, image):
         return self.model(image)
 
-    def on_train_epoch_end(self, unused = None):
+    def on_train_epoch_end(self, unused=None):
         for m in self.train_metrics:
             m.reset()