Repository: mlcommons/GaNDLF

Issue:
Need to switch from BCE
**Describe the bug**
A clear and concise description of what the bug is.
When running GaNDLF for a 3D U-Net segmentation task with a weighted `dcce` loss function, the following error occurred. 

```python-traceback
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [31,0,0], thread: [0,0,0] Assertion `input_val >= zero && input_val <= one` failed.

Traceback (most recent call last):
  File "/cbica/home/patis/comp_space/testing/gandlf_mine/gandlf_run", line 157, in <module>
    main()
  File "/cbica/home/patis/comp_space/testing/gandlf_mine/gandlf_run", line 129, in main
    reset_prev=reset_prev,
  File "/gpfs/fs001/cbica/comp_space/patis/testing/gandlf_mine/GANDLF/training_manager.py", line 319, in TrainingManager_split
    testing_data=None,
  File "/gpfs/fs001/cbica/comp_space/patis/testing/gandlf_mine/GANDLF/training_loop.py", line 792, in training_loop
    model, val_dataloader, scheduler, params, epoch, mode="validation"
  File "/gpfs/fs001/cbica/comp_space/patis/testing/gandlf_mine/GANDLF/training_loop.py", line 473, in validate_network
    result = step(model, image, label, params)
  File "/gpfs/fs001/cbica/comp_space/patis/testing/gandlf_mine/GANDLF/training_loop.py", line 94, in step
    loss, metric_output = get_loss_and_metrics(label, output, params)
  File "/gpfs/fs001/cbica/comp_space/patis/testing/gandlf_mine/GANDLF/parameterParsing.py", line 482, in get_loss_and_metrics
    metric_function(predicted, ground_truth, params).cpu().data.item()
RuntimeError: CUDA error: device-side assert triggered
```

This indicated that there was a problem with the use of the `BCELoss()` function in line 79 of `losses.py`:
```python
loss = torch.nn.BCELoss()
```

After further investigation, it seems that changing line 79 to use the `BCEWithLogitsLoss()` function solves the issue:
```python
   loss = torch.nn.BCEWithLogitsLoss()
```


**To Reproduce**
Steps to reproduce the behavior:
Run this command:
```
source activate /gpfs/fs001/cbica/home/ahluwalv/dbt_deep_learning/GaNDLF/venv

/gpfs/fs001/cbica/home/ahluwalv/dbt_deep_learning/GaNDLF/venv/bin/python /cbica/home/ahluwalv/dbt_deep_learning/GaNDLF/gandlf_run -config /cbica/home/patis/comp_space/testing/gandlf_mine/exp_vinny/A/config.yaml -data /cbica/home/ahluwalv/dbt_deep_learning/train_dbt.csv,/cbica/home/ahluwalv/dbt_deep_learning/val_dbt.csv -output /cbica/home/ahluwalv/dbt_deep_learning/train_logs_augmentation/ -train 1 -device cuda  -reset_prev True
```

**Expected behavior**
It is expected that when using the loss function, that values remain between `[0,1]`.



**GaNDLF Version**
0.0.10-dev

**Desktop (please complete the following information):**
 - OS: Linux

**Additional context**
This issue occurred during epoch 4, so its occurrence may be random. Ensuring it doesn't happen again will probably be accomplished by the suggested fix.



Discussion:

Resolving patch:
diff --git a/GANDLF/losses.py b/GANDLF/losses.py
--- a/GANDLF/losses.py
+++ b/GANDLF/losses.py
@@ -10,9 +10,8 @@ def dice(inp, target):
     iflat = inp.contiguous().view(-1)
     tflat = target.contiguous().view(-1)
     intersection = (iflat * tflat).sum()
-    return (2.0 * intersection + smooth) / (
-        iflat.sum() + tflat.sum() + smooth
-    )  # 2 * intersection / union
+    # 2 * intersection / union
+    return (2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth)
 
 
 def MCD(pm, gt, num_class, weights=None, ignore_class=None, loss_type=0):
@@ -35,9 +34,8 @@ def MCD(pm, gt, num_class, weights=None, ignore_class=None, loss_type=0):
             if loss_type == 1:
                 currentDice = 1 - currentDice  # subtract from 1 because this is a loss
             elif loss_type == 2:
-                currentDice = -torch.log(
-                    currentDice + torch.finfo(torch.float32).eps
-                )  # negative because we want positive losses
+                # negative because we want positive losses
+                currentDice = -torch.log(currentDice + torch.finfo(torch.float32).eps)
             if weights is not None:
                 currentDice = currentDice * weights[i]
             acc_dice += currentDice
@@ -80,6 +78,7 @@ def CE_Logits(out, target):
     loss_val = loss(iflat, tflat)
     return loss_val
 
+
 def CE(out, target):
     iflat = out.contiguous().view(-1)
     tflat = target.contiguous().view(-1)
@@ -87,24 +86,23 @@ def CE(out, target):
     loss_val = loss(iflat, tflat)
     return loss_val
 
-def CCE_Logits(out, target, params):
-    acc_ce_loss = 0
-    target = one_hot(target, params["model"]["class_list"]).type(out.dtype)
-    for i in range(0, len(params["model"]["class_list"])):
-        curr_ce_loss = CE_Logits(out[:, i, ...], target[:, i, ...])
-        if params["weights"] is not None:
-            curr_ce_loss = curr_ce_loss * params["weights"][i]
-        acc_ce_loss += curr_ce_loss
-    if params["weights"] is None:
-        acc_ce_loss /= len(params["model"]["class_list"])
-    return acc_ce_loss
 
-# This is wrong, that is not how categorical cross entropy works
-def CCE(out, target, params):
+def CCE_Generic(out, target, params, CCE_Type):
+    """Generic function to calculate CCE loss
+
+    Args:
+        out (torch.tensor): The predicted output value for each pixel. dimension: [batch, class, x, y, z].
+        target (torch.tensor): The ground truth label for each pixel. dimension: [batch, class, x, y, z] factorial_class_list.
+        params (dict): The parameter dictionary.
+        CCE_Type (torch.nn): The CE loss function type.
+
+    Returns:
+        torch.tensor: The final loss value after taking multiple classes into consideration
+    """
     acc_ce_loss = 0
     target = one_hot(target, params["model"]["class_list"]).type(out.dtype)
     for i in range(0, len(params["model"]["class_list"])):
-        curr_ce_loss = CE(out[:, i, ...], target[:, i, ...])
+        curr_ce_loss = CCE_Type(out[:, i, ...], target[:, i, ...])
         if params["weights"] is not None:
             curr_ce_loss = curr_ce_loss * params["weights"][i]
         acc_ce_loss += curr_ce_loss
@@ -114,13 +112,13 @@ def CCE(out, target, params):
 
 
 def DCCE(pm, gt, params):
-    dcce_loss = MCD_loss(pm, gt, params) + CCE(pm, gt, params)
+    dcce_loss = MCD_loss(pm, gt, params) + CCE_Generic(pm, gt, params, CE)
     return dcce_loss
-    
+
+
 def DCCE_Logits(pm, gt, params):
-    dcce_loss = MCD_loss(pm, gt, params) + CCE_Logits(pm, gt, params)
+    dcce_loss = MCD_loss(pm, gt, params) + CCE_Generic(pm, gt, params, CE_Logits)
     return dcce_loss
-    
 
 
 def tversky(inp, target, alpha):
diff --git a/GANDLF/parseConfig.py b/GANDLF/parseConfig.py
--- a/GANDLF/parseConfig.py
+++ b/GANDLF/parseConfig.py
@@ -146,9 +146,8 @@ def parseConfig(config_file_path, version_check=True):
                             params["loss_function"][key] = {}
                             params["loss_function"][key]["reduction"] = "mean"
                     else:
-                        params[
-                            "loss_function"
-                        ] = key  # use simple string for other functions - can be extended with parameters, if needed
+                        # use simple string for other functions - can be extended with parameters, if needed
+                        params["loss_function"] = key
             else:
                 defineDefaultLoss = True
         else:


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -431,7 +431,7 @@ def test_clip_train_classification_rad_3d(device):
 
 def test_normtype_train_segmentation_rad_3d(device):
     # read and initialize parameters for specific data dimension
-    print("Starting 3D Rad segmentation tests")
+    print("Starting 3D Rad segmentation tests for normtype")
     # read and parse csv
     # read and initialize parameters for specific data dimension
     parameters = parseConfig(
@@ -466,7 +466,7 @@ def test_normtype_train_segmentation_rad_3d(device):
 
 
 def test_metrics_segmentation_rad_2d(device):
-    print("Starting 2D Rad segmentation tests")
+    print("Starting 2D Rad segmentation tests for metrics")
     # read and parse csv
     parameters = parseConfig(
         testingDir + "/config_segmentation.yaml", version_check=False
@@ -493,3 +493,37 @@ def test_metrics_segmentation_rad_2d(device):
     shutil.rmtree(outputDir)  # overwrite previous results
 
     print("passed")
+
+
+def test_losses_segmentation_rad_2d(device):
+    print("Starting 2D Rad segmentation tests for losses")
+    # read and parse csv
+    parameters = parseConfig(
+        testingDir + "/config_segmentation.yaml", version_check=False
+    )
+    training_data, parameters["headers"] = parseTrainingCSV(
+        inputDir + "/train_2d_rad_segmentation.csv"
+    )
+    parameters = populate_header_in_parameters(parameters, parameters["headers"])
+    parameters["patch_size"] = patch_size["2D"]
+    parameters["model"]["dimension"] = 2
+    parameters["model"]["class_list"] = [0, 255]
+    # disabling amp because some losses do not support Half, yet 
+    parameters["model"]["amp"] = False
+    parameters["model"]["num_channels"] = 3
+    parameters["model"]["architecture"] = "unet"
+    parameters["metrics"] = ["dice"]
+    # loop through selected models and train for single epoch
+    for loss_type in ["dc", "dc_log", "dcce", "dcce_logits"]:
+        parameters["loss_function"] = loss_type
+        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        TrainingManager(
+            dataframe=training_data,
+            outputDir=outputDir,
+            parameters=parameters,
+            device=device,
+            reset_prev=True,
+        )
+        shutil.rmtree(outputDir)  # overwrite previous results
+    print("passed")
+