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
old mode 100644
new mode 100755
--- a/GANDLF/losses.py
+++ b/GANDLF/losses.py
@@ -73,6 +73,13 @@ def MCD_log_loss(pm, gt, params):
     return MCD(pm, gt, len(params["model"]["class_list"]), params["weights"], None, 2)
 
 
+def CE_Logits(out, target):
+    iflat = out.contiguous().view(-1)
+    tflat = target.contiguous().view(-1)
+    loss = torch.nn.BCEWithLogitsLoss()
+    loss_val = loss(iflat, tflat)
+    return loss_val
+
 def CE(out, target):
     iflat = out.contiguous().view(-1)
     tflat = target.contiguous().view(-1)
@@ -80,6 +87,17 @@ 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):
@@ -98,6 +116,11 @@ def CCE(out, target, params):
 def DCCE(pm, gt, params):
     dcce_loss = MCD_loss(pm, gt, params) + CCE(pm, gt, params)
     return dcce_loss
+    
+def DCCE_Logits(pm, gt, params):
+    dcce_loss = MCD_loss(pm, gt, params) + CCE_Logits(pm, gt, params)
+    return dcce_loss
+    
 
 
 def tversky(inp, target, alpha):
@@ -194,6 +217,8 @@ def fetch_loss_function(loss_name, params):
         loss_function = MCD_log_loss
     elif loss_name == "dcce":
         loss_function = DCCE
+    elif loss_name == "dcce_logits":
+        loss_function = DCCE_Logits
     elif loss_name == "ce":
         loss_function = CE
     elif loss_name == "mse":


Test patch: