focal

models/deepcrack_model.py

@@ -4,7 +4,7 @@
 import numpy as np
 import itertools
 from .base_model import BaseModel
-from .deepcrack_networks import define_deepcrack
+from .deepcrack_networks import define_deepcrack, BinaryFocalLoss
 
 class DeepCrackModel(BaseModel):
     """
@@ -50,7 +50,10 @@ def __init__(self, opt):
             # define loss functions
             #self.weight = torch.from_numpy(np.array([0.0300, 1.0000], dtype='float32')).float().to(self.device)
             #self.criterionSeg = torch.nn.CrossEntropyLoss(weight=self.weight)
-            self.criterionSeg = nn.BCEWithLogitsLoss(size_average=True, reduce=True, 
+            if self.opt.loss_mode == 'focal':
+                self.criterionSeg = BinaryFocalLoss()
+            else: 
+                self.criterionSeg = nn.BCEWithLogitsLoss(size_average=True, reduce=True, 
                     pos_weight=torch.tensor(1.0/3e-2).to(self.device))
             self.weight_side = [0.5, 0.75, 1.0, 0.75, 0.5]
 

models/deepcrack_networks.py

@@ -88,3 +88,23 @@ def define_deepcrack(in_nc,
                      gpu_ids=[]):
     net = DeepCrackNet(in_nc, num_classes, ngf, norm)
     return init_net(net, init_type, init_gain, gpu_ids)
+
+
+class BinaryFocalLoss(nn.Module):
+    def __init__(self, alpha=1, gamma=2, logits=False, size_average=True):
+        super(BinaryFocalLoss, self).__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.logits = logits
+        self.size_average = size_average
+        self.criterion = nn.BCEWithLogitsLoss(reduction='none')
+
+    def forward(self, inputs, targets):
+        BCE_loss = self.criterion(inputs, targets)
+        pt = torch.exp(-BCE_loss)
+        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss
+
+        if self.size_average:
+            return F_loss.mean()
+        else:
+            return F_loss.sum()

options/base_options.py

@@ -64,6 +64,7 @@ def initialize(self, parser):
         #parser.add_argument('--use_l1', type=int, default=1, help='using l1 loss')
 
         parser.add_argument('--use_selu', type=int, default=1, help='using selu active function')
+        parser.add_argument('--loss_mode', type=str, default='focal', help='[bce | focal]')
         self.initialized = True
         return parser
 

scripts/train_deepcrack.sh

@@ -9,6 +9,7 @@ BATCH_SIZE=1
 NUM_CLASSES=1
 LOAD_WIDTH=256
 LOAD_HEIGHT=256
+LOSS_MODE=focal
 
 NORM=batch
 NITER=400
@@ -30,4 +31,5 @@ python3 train.py \
   --load_width ${LOAD_WIDTH} \
   --load_height ${LOAD_HEIGHT} \
   --no_flip 0 \
-  --display_id 0
+  --display_id 0 \
+  --loss_mode ${LOSS_MODE}