implemented 2d dice score and pixel normalization. Training loss is actually pretty good now

Author: harshakmohan

tutorial/train.py

@@ -7,6 +7,7 @@
 import torch.optim as optim
 from model import UNET
 from utils import (load_checkpoint, save_checkpoint, get_loaders, check_accuracy, save_predictions_as_imgs)
+from utils import DiceLoss2D
 
 # Hyperparameters
 LEARNING_RATE = 1e-4
@@ -47,7 +48,8 @@ def train_fn(loader, model, optimizer, loss_fn, scaler):
 def main():
 
     model = UNET(in_channels=3, out_channels=1).to(device=DEVICE)
-    loss_fn = nn.BCEWithLogitsLoss() # LOSS FUNCTION DEFINED HERE. Perhaps change it to dice score
+    #loss_fn = nn.BCEWithLogitsLoss() # LOSS FUNCTION DEFINED HERE. Perhaps change it to dice score
+    loss_fn = DiceLoss2D()
     optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
 
     train_loader, val_loader = get_loaders(train_dir=TRAIN_IMG_DIR,

tutorial/utils.py

@@ -114,7 +114,6 @@ class DiceLoss2D(nn.Module):
 
     def __init__(self, skip_bg=False):
         super(DiceLoss2D, self).__init__()
-
         self.skip_bg = skip_bg
 
     def forward(self, inputs, target):