We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Penalty weights are incorrect
**Describe the bug**
See title.

**To Reproduce**
Steps to reproduce the behavior:
1. Any segmentation/classification problem
2. Take a print at the end of `GANDLF.utils.tensor.get_class_imbalance_weights`
3. See that `penalty_dict` and `weights_dict` are the same

For example, for a small I-SPY dataset:
```
penalty_dict: {0: 0.9928165871714076, 1: 0.0071834128285923945} 
weights_dict: {0: 0.9928165871714076, 1: 0.0071834128285923945}
```

**Expected behavior**
They should be `1-each_other`. In the above example, it should be this:

```
penalty_dict: {0: 0.0071834128285923945, 1: 0.9928165871714076} 
weights_dict: {0: 0.9928165871714076,    1: 0.0071834128285923945}
```

**Screenshots**
N.A.

**GaNDLF Version**
0.0.12-dev

**Desktop (please complete the following information):**
N.A.

**Additional context**
Reported by @iskobleva 


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `GANDLF/utils/tensor.py`
Content:
```
1 import os, sys
2 import numpy as np
3 import torch
4 import torch.nn as nn
5 import torchio
6 
7 
8 def one_hot(segmask_array, class_list):
9     """
10     This function creates a one-hot-encoded mask from the segmentation mask Tensor and specified class list
11 
12     Args:
13         segmask_array (torch.Tensor): The segmentation mask Tensor.
14         class_list (list): The list of classes based on which one-hot encoding needs to happen.
15 
16     Returns:
17         torch.Tensor: The one-hot encoded torch.Tensor
18     """
19     batch_size = segmask_array.shape[0]
20     batch_stack = []
21     for b in range(batch_size):
22         one_hot_stack = []
23         segmask_array_iter = segmask_array[b, 0]
24         bin_mask = segmask_array_iter == 0  # initialize bin_mask
25         for (
26             _class
27         ) in class_list:  # this implementation allows users to combine logical operands
28             if isinstance(_class, str):
29                 if "||" in _class:  # special case
30                     class_split = _class.split("||")
31                     bin_mask = segmask_array_iter == int(class_split[0])
32                     for i in range(1, len(class_split)):
33                         bin_mask = bin_mask | (
34                             segmask_array_iter == int(class_split[i])
35                         )
36                 elif "|" in _class:  # special case
37                     class_split = _class.split("|")
38                     bin_mask = segmask_array_iter == int(class_split[0])
39                     for i in range(1, len(class_split)):
40                         bin_mask = bin_mask | (
41                             segmask_array_iter == int(class_split[i])
42                         )
43                 else:
44                     # assume that it is a simple int
45                     bin_mask = segmask_array_iter == int(_class)
46             else:
47                 bin_mask = segmask_array_iter == int(_class)
48                 bin_mask = bin_mask.long()
49             one_hot_stack.append(bin_mask)
50         one_hot_stack = torch.stack(one_hot_stack)
51         batch_stack.append(one_hot_stack)
52     batch_stack = torch.stack(batch_stack)
53     return batch_stack
54 
55 
56 def reverse_one_hot(predmask_array, class_list):
57     """
58     This function creates a full segmentation mask Tensor from a one-hot-encoded mask and specified class list
59 
60     Args:
61         predmask_array (torch.Tensor): The predicted segmentation mask Tensor.
62         class_list (list): The list of classes based on which one-hot encoding needs to happen.
63 
64     Returns:
65         torch.Tensor: The final mask torch.Tensor.
66     """
67     if isinstance(predmask_array, torch.Tensor):
68         array_to_consider = predmask_array.cpu().numpy()
69     else:
70         array_to_consider = predmask_array
71     idx_argmax = np.argmax(array_to_consider, axis=0)
72     final_mask = 0
73     special_cases_to_check = ["||"]
74     special_case_detected = False
75     max_current = 0
76 
77     for _class in class_list:
78         for case in special_cases_to_check:
79             if isinstance(_class, str):
80                 if case in _class:  # check if any of the special cases are present
81                     special_case_detected = True
82                     class_split = _class.split(
83                         case
84                     )  # if present, then split the sub-class
85                     for i in class_split:  # find the max for computation later on
86                         if int(i) > max_current:
87                             max_current = int(i)
88 
89     if special_case_detected:
90         start_idx = 0
91         if (class_list[0] == 0) or (class_list[0] == "0"):
92             start_idx = 1
93 
94         final_mask = np.asarray(
95             predmask_array[start_idx, :, :, :], dtype=int
96         )  # predmask_array[0,:,:,:].long()
97         start_idx += 1
98         for i in range(start_idx, len(class_list)):
99             final_mask += np.asarray(
100                 predmask_array[0, :, :, :], dtype=int
101             )  # predmask_array[i,:,:,:].long()
102             # temp_sum = torch.sum(output)
103         # output_2 = (max_current - torch.sum(output)) % max_current
104         # test_2 = 1
105     else:
106         for idx, _class in enumerate(class_list):
107             final_mask = final_mask + (idx_argmax == idx) * _class
108     return final_mask
109 
110 
111 def send_model_to_device(model, amp, device, optimizer):
112     """
113     This function reads the environment variable(s) and send model to correct device
114 
115     Args:
116         model (torch.nn.Module): The model that needs to be sent to specified device.
117         amp (bool): Whether automatic mixed precision is to be used.
118         device (str): Device type.
119         optimizer (torch.optim): The optimizer for training.
120 
121     Returns:
122         torch.nn.Module: The model after it has been sent to specified device
123         bool: Whether automatic mixed precision is to be used or not.
124         torch.device: Device type.
125     """
126     if device != "cpu":
127         if os.environ.get("CUDA_VISIBLE_DEVICES") is None:
128             sys.exit(
129                 "Please set the environment variable 'CUDA_VISIBLE_DEVICES' correctly before trying to run GANDLF on GPU"
130             )
131 
132         dev = os.environ.get("CUDA_VISIBLE_DEVICES")
133         # multi-gpu support
134         # ###
135         # # https://discuss.pytorch.org/t/cuda-visible-devices-make-gpu-disappear/21439/17?u=sarthakpati
136         # ###
137         if "," in dev:
138             device = torch.device("cuda")
139             model = nn.DataParallel(model, "[" + dev + "]")
140         else:
141             print("Device requested via CUDA_VISIBLE_DEVICES: ", dev)
142             print("Total number of CUDA devices: ", torch.cuda.device_count())
143 
144             # if only a single visible device, it will be indexed as '0'
145             if torch.cuda.device_count() == 1:
146                 dev = "0"
147 
148             dev_int = int(dev)
149             print("Device finally used: ", dev)
150             # device = torch.device('cuda:' + dev)
151             device = torch.device("cuda")
152             print("Sending model to aforementioned device")
153             model = model.to(device)
154             print(
155                 "Memory Total : ",
156                 round(
157                     torch.cuda.get_device_properties(dev_int).total_memory / 1024 ** 3,
158                     1,
159                 ),
160                 "GB, Allocated: ",
161                 round(torch.cuda.memory_allocated(dev_int) / 1024 ** 3, 1),
162                 "GB, Cached: ",
163                 round(torch.cuda.memory_reserved(dev_int) / 1024 ** 3, 1),
164                 "GB",
165             )
166 
167         print(
168             "Device - Current: %s Count: %d Name: %s Availability: %s"
169             % (
170                 torch.cuda.current_device(),
171                 torch.cuda.device_count(),
172                 torch.cuda.get_device_name(device),
173                 torch.cuda.is_available(),
174             )
175         )
176 
177         if not (optimizer is None):
178             # ensuring optimizer is in correct device - https://github.com/pytorch/pytorch/issues/8741
179             optimizer.load_state_dict(optimizer.state_dict())
180 
181     else:
182         dev = -1
183         device = torch.device("cpu")
184         model.cpu()
185         amp = False
186         print("Since Device is CPU, Mixed Precision Training is set to False")
187 
188     return model, amp, device
189 
190 
191 def get_class_imbalance_weights(training_data_loader, parameters):
192     """
193     This function calculates the penalty that is used for validation loss in multi-class problems
194 
195     Args:
196         training_data_loader (torch.utils.data.DataLoader): The training data loader.
197         parameters (dict): The parameters passed by the user yaml.
198 
199     Returns:
200         dict: The penalty weights for different classes under consideration.
201     """
202     abs_dict = {}  # absolute counts for each class
203     weights_dict = {}  # average for "weighted averaging"
204     penalty_dict = None  # penalty for misclassification
205     # basically, do this for segmentation/classification tasks
206 
207     if parameters["problem_type"] != "regression":
208         penalty_dict = {}
209         for i in range(0, len(parameters["model"]["class_list"])):
210             abs_dict[i] = 0
211             penalty_dict[i] = 0
212 
213     penalty_loader = training_data_loader
214 
215     # get the weights for use for dice loss
216     total_counter = 0
217 
218     # For regression dice penalty need not be taken account
219     # For classification this should be calculated on the basis of predicted labels and mask
220     # iterate through full penalty data
221     for _, (subject) in enumerate(penalty_loader):
222 
223         # segmentation needs masks to be one-hot encoded
224         if parameters["problem_type"] == "segmentation":
225             # accumulate dice weights for each label
226             mask = subject["label"][torchio.DATA]
227             one_hot_mask = one_hot(mask, parameters["model"]["class_list"])
228             for i in range(0, len(parameters["model"]["class_list"])):
229                 currentNumber = torch.nonzero(
230                     one_hot_mask[:, i, :, :, :], as_tuple=False
231                 ).size(0)
232                 # class-specific non-zero voxels
233                 abs_dict[i] += currentNumber
234                 # total number of non-zero voxels to be considered
235                 total_counter += currentNumber
236 
237         # for classification, the value needs to be used directly
238         elif parameters["problem_type"] == "classification":
239             # accumulate weights for each label
240             value_to_predict = subject["value_0"][0]
241             for i in range(0, len(parameters["model"]["class_list"])):
242                 if value_to_predict == i:
243                     abs_dict[i] += 1
244                     # we only want to increase the counter for those subjects that are defined in the class_list
245                     total_counter += 1
246 
247     # Normalize class weights
248     weights_dict = {key: val / total_counter for key, val in abs_dict.items()}
249 
250     # get the penalty values - abs_dict contains the overall number for each class in the penalty data
251     for i in range(0, len(parameters["model"]["class_list"])):
252         penalty = total_counter  # start with the assumption that all the non-zero voxels (segmentation) or activate labels (classification) make up the penalty
253         for j in range(0, len(parameters["model"]["class_list"])):
254             if i != j:  # for differing classes, subtract the current weight
255                 penalty -= abs_dict[j]
256 
257         # finally, the "penalty" variable contains the total number of voxels/activations that are not part of the current class
258         # this is to be used to weight the loss function
259         # adding epsilon to avoid division by zero
260         penalty_dict[i] = (penalty + sys.float_info.epsilon) / (
261             total_counter + sys.float_info.epsilon
262         )
263 
264     return penalty_dict, weights_dict
265
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/GANDLF/utils/tensor.py b/GANDLF/utils/tensor.py
--- a/GANDLF/utils/tensor.py
+++ b/GANDLF/utils/tensor.py
@@ -22,9 +22,8 @@
         one_hot_stack = []
         segmask_array_iter = segmask_array[b, 0]
         bin_mask = segmask_array_iter == 0  # initialize bin_mask
-        for (
-            _class
-        ) in class_list:  # this implementation allows users to combine logical operands
+        # this implementation allows users to combine logical operands
+        for _class in class_list:
             if isinstance(_class, str):
                 if "||" in _class:  # special case
                     class_split = _class.split("||")
@@ -79,9 +78,8 @@
             if isinstance(_class, str):
                 if case in _class:  # check if any of the special cases are present
                     special_case_detected = True
-                    class_split = _class.split(
-                        case
-                    )  # if present, then split the sub-class
+                    # if present, then split the sub-class
+                    class_split = _class.split(case)
                     for i in class_split:  # find the max for computation later on
                         if int(i) > max_current:
                             max_current = int(i)
@@ -91,9 +89,7 @@
         if (class_list[0] == 0) or (class_list[0] == "0"):
             start_idx = 1
 
-        final_mask = np.asarray(
-            predmask_array[start_idx, :, :, :], dtype=int
-        )  # predmask_array[0,:,:,:].long()
+        final_mask = np.asarray(predmask_array[start_idx, :, :, :], dtype=int)
         start_idx += 1
         for i in range(start_idx, len(class_list)):
             final_mask += np.asarray(
@@ -245,20 +241,21 @@
                     total_counter += 1
 
     # Normalize class weights
-    weights_dict = {key: val / total_counter for key, val in abs_dict.items()}
-
-    # get the penalty values - abs_dict contains the overall number for each class in the penalty data
-    for i in range(0, len(parameters["model"]["class_list"])):
-        penalty = total_counter  # start with the assumption that all the non-zero voxels (segmentation) or activate labels (classification) make up the penalty
-        for j in range(0, len(parameters["model"]["class_list"])):
-            if i != j:  # for differing classes, subtract the current weight
-                penalty -= abs_dict[j]
-
-        # finally, the "penalty" variable contains the total number of voxels/activations that are not part of the current class
-        # this is to be used to weight the loss function
-        # adding epsilon to avoid division by zero
-        penalty_dict[i] = (penalty + sys.float_info.epsilon) / (
-            total_counter + sys.float_info.epsilon
-        )
+    weights_dict = {
+        key: (val + sys.float_info.epsilon) / total_counter
+        for key, val in abs_dict.items()
+    }
+
+    # get the raw penalty values
+    penalty = {
+        key: total_counter / (len(abs_dict) * (val + sys.float_info.epsilon))
+        for key, val in abs_dict.items()
+    }
+    # normalize penalty to sum of 1
+    penalty_sum = np.fromiter(penalty.values(), dtype=np.float64).sum()
+    penalty_dict = {
+        key: (val + sys.float_info.epsilon) / penalty_sum
+        for key, val in penalty.items()
+    }
 
     return penalty_dict, weights_dict

{"golden_diff": "diff --git a/GANDLF/utils/tensor.py b/GANDLF/utils/tensor.py\n--- a/GANDLF/utils/tensor.py\n+++ b/GANDLF/utils/tensor.py\n@@ -22,9 +22,8 @@\n         one_hot_stack = []\n         segmask_array_iter = segmask_array[b, 0]\n         bin_mask = segmask_array_iter == 0  # initialize bin_mask\n-        for (\n-            _class\n-        ) in class_list:  # this implementation allows users to combine logical operands\n+        # this implementation allows users to combine logical operands\n+        for _class in class_list:\n             if isinstance(_class, str):\n                 if \"||\" in _class:  # special case\n                     class_split = _class.split(\"||\")\n@@ -79,9 +78,8 @@\n             if isinstance(_class, str):\n                 if case in _class:  # check if any of the special cases are present\n                     special_case_detected = True\n-                    class_split = _class.split(\n-                        case\n-                    )  # if present, then split the sub-class\n+                    # if present, then split the sub-class\n+                    class_split = _class.split(case)\n                     for i in class_split:  # find the max for computation later on\n                         if int(i) > max_current:\n                             max_current = int(i)\n@@ -91,9 +89,7 @@\n         if (class_list[0] == 0) or (class_list[0] == \"0\"):\n             start_idx = 1\n \n-        final_mask = np.asarray(\n-            predmask_array[start_idx, :, :, :], dtype=int\n-        )  # predmask_array[0,:,:,:].long()\n+        final_mask = np.asarray(predmask_array[start_idx, :, :, :], dtype=int)\n         start_idx += 1\n         for i in range(start_idx, len(class_list)):\n             final_mask += np.asarray(\n@@ -245,20 +241,21 @@\n                     total_counter += 1\n \n     # Normalize class weights\n-    weights_dict = {key: val / total_counter for key, val in abs_dict.items()}\n-\n-    # get the penalty values - abs_dict contains the overall number for each class in the penalty data\n-    for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n-        penalty = total_counter  # start with the assumption that all the non-zero voxels (segmentation) or activate labels (classification) make up the penalty\n-        for j in range(0, len(parameters[\"model\"][\"class_list\"])):\n-            if i != j:  # for differing classes, subtract the current weight\n-                penalty -= abs_dict[j]\n-\n-        # finally, the \"penalty\" variable contains the total number of voxels/activations that are not part of the current class\n-        # this is to be used to weight the loss function\n-        # adding epsilon to avoid division by zero\n-        penalty_dict[i] = (penalty + sys.float_info.epsilon) / (\n-            total_counter + sys.float_info.epsilon\n-        )\n+    weights_dict = {\n+        key: (val + sys.float_info.epsilon) / total_counter\n+        for key, val in abs_dict.items()\n+    }\n+\n+    # get the raw penalty values\n+    penalty = {\n+        key: total_counter / (len(abs_dict) * (val + sys.float_info.epsilon))\n+        for key, val in abs_dict.items()\n+    }\n+    # normalize penalty to sum of 1\n+    penalty_sum = np.fromiter(penalty.values(), dtype=np.float64).sum()\n+    penalty_dict = {\n+        key: (val + sys.float_info.epsilon) / penalty_sum\n+        for key, val in penalty.items()\n+    }\n \n     return penalty_dict, weights_dict\n", "issue": "Penalty weights are incorrect\n**Describe the bug**\r\nSee title.\r\n\r\n**To Reproduce**\r\nSteps to reproduce the behavior:\r\n1. Any segmentation/classification problem\r\n2. Take a print at the end of `GANDLF.utils.tensor.get_class_imbalance_weights`\r\n3. See that `penalty_dict` and `weights_dict` are the same\r\n\r\nFor example, for a small I-SPY dataset:\r\n```\r\npenalty_dict: {0: 0.9928165871714076, 1: 0.0071834128285923945} \r\nweights_dict: {0: 0.9928165871714076, 1: 0.0071834128285923945}\r\n```\r\n\r\n**Expected behavior**\r\nThey should be `1-each_other`. In the above example, it should be this:\r\n\r\n```\r\npenalty_dict: {0: 0.0071834128285923945, 1: 0.9928165871714076} \r\nweights_dict: {0: 0.9928165871714076,    1: 0.0071834128285923945}\r\n```\r\n\r\n**Screenshots**\r\nN.A.\r\n\r\n**GaNDLF Version**\r\n0.0.12-dev\r\n\r\n**Desktop (please complete the following information):**\r\nN.A.\r\n\r\n**Additional context**\r\nReported by @iskobleva \r\n\n", "before_files": [{"content": "import os, sys\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torchio\n\n\ndef one_hot(segmask_array, class_list):\n    \"\"\"\n    This function creates a one-hot-encoded mask from the segmentation mask Tensor and specified class list\n\n    Args:\n        segmask_array (torch.Tensor): The segmentation mask Tensor.\n        class_list (list): The list of classes based on which one-hot encoding needs to happen.\n\n    Returns:\n        torch.Tensor: The one-hot encoded torch.Tensor\n    \"\"\"\n    batch_size = segmask_array.shape[0]\n    batch_stack = []\n    for b in range(batch_size):\n        one_hot_stack = []\n        segmask_array_iter = segmask_array[b, 0]\n        bin_mask = segmask_array_iter == 0  # initialize bin_mask\n        for (\n            _class\n        ) in class_list:  # this implementation allows users to combine logical operands\n            if isinstance(_class, str):\n                if \"||\" in _class:  # special case\n                    class_split = _class.split(\"||\")\n                    bin_mask = segmask_array_iter == int(class_split[0])\n                    for i in range(1, len(class_split)):\n                        bin_mask = bin_mask | (\n                            segmask_array_iter == int(class_split[i])\n                        )\n                elif \"|\" in _class:  # special case\n                    class_split = _class.split(\"|\")\n                    bin_mask = segmask_array_iter == int(class_split[0])\n                    for i in range(1, len(class_split)):\n                        bin_mask = bin_mask | (\n                            segmask_array_iter == int(class_split[i])\n                        )\n                else:\n                    # assume that it is a simple int\n                    bin_mask = segmask_array_iter == int(_class)\n            else:\n                bin_mask = segmask_array_iter == int(_class)\n                bin_mask = bin_mask.long()\n            one_hot_stack.append(bin_mask)\n        one_hot_stack = torch.stack(one_hot_stack)\n        batch_stack.append(one_hot_stack)\n    batch_stack = torch.stack(batch_stack)\n    return batch_stack\n\n\ndef reverse_one_hot(predmask_array, class_list):\n    \"\"\"\n    This function creates a full segmentation mask Tensor from a one-hot-encoded mask and specified class list\n\n    Args:\n        predmask_array (torch.Tensor): The predicted segmentation mask Tensor.\n        class_list (list): The list of classes based on which one-hot encoding needs to happen.\n\n    Returns:\n        torch.Tensor: The final mask torch.Tensor.\n    \"\"\"\n    if isinstance(predmask_array, torch.Tensor):\n        array_to_consider = predmask_array.cpu().numpy()\n    else:\n        array_to_consider = predmask_array\n    idx_argmax = np.argmax(array_to_consider, axis=0)\n    final_mask = 0\n    special_cases_to_check = [\"||\"]\n    special_case_detected = False\n    max_current = 0\n\n    for _class in class_list:\n        for case in special_cases_to_check:\n            if isinstance(_class, str):\n                if case in _class:  # check if any of the special cases are present\n                    special_case_detected = True\n                    class_split = _class.split(\n                        case\n                    )  # if present, then split the sub-class\n                    for i in class_split:  # find the max for computation later on\n                        if int(i) > max_current:\n                            max_current = int(i)\n\n    if special_case_detected:\n        start_idx = 0\n        if (class_list[0] == 0) or (class_list[0] == \"0\"):\n            start_idx = 1\n\n        final_mask = np.asarray(\n            predmask_array[start_idx, :, :, :], dtype=int\n        )  # predmask_array[0,:,:,:].long()\n        start_idx += 1\n        for i in range(start_idx, len(class_list)):\n            final_mask += np.asarray(\n                predmask_array[0, :, :, :], dtype=int\n            )  # predmask_array[i,:,:,:].long()\n            # temp_sum = torch.sum(output)\n        # output_2 = (max_current - torch.sum(output)) % max_current\n        # test_2 = 1\n    else:\n        for idx, _class in enumerate(class_list):\n            final_mask = final_mask + (idx_argmax == idx) * _class\n    return final_mask\n\n\ndef send_model_to_device(model, amp, device, optimizer):\n    \"\"\"\n    This function reads the environment variable(s) and send model to correct device\n\n    Args:\n        model (torch.nn.Module): The model that needs to be sent to specified device.\n        amp (bool): Whether automatic mixed precision is to be used.\n        device (str): Device type.\n        optimizer (torch.optim): The optimizer for training.\n\n    Returns:\n        torch.nn.Module: The model after it has been sent to specified device\n        bool: Whether automatic mixed precision is to be used or not.\n        torch.device: Device type.\n    \"\"\"\n    if device != \"cpu\":\n        if os.environ.get(\"CUDA_VISIBLE_DEVICES\") is None:\n            sys.exit(\n                \"Please set the environment variable 'CUDA_VISIBLE_DEVICES' correctly before trying to run GANDLF on GPU\"\n            )\n\n        dev = os.environ.get(\"CUDA_VISIBLE_DEVICES\")\n        # multi-gpu support\n        # ###\n        # # https://discuss.pytorch.org/t/cuda-visible-devices-make-gpu-disappear/21439/17?u=sarthakpati\n        # ###\n        if \",\" in dev:\n            device = torch.device(\"cuda\")\n            model = nn.DataParallel(model, \"[\" + dev + \"]\")\n        else:\n            print(\"Device requested via CUDA_VISIBLE_DEVICES: \", dev)\n            print(\"Total number of CUDA devices: \", torch.cuda.device_count())\n\n            # if only a single visible device, it will be indexed as '0'\n            if torch.cuda.device_count() == 1:\n                dev = \"0\"\n\n            dev_int = int(dev)\n            print(\"Device finally used: \", dev)\n            # device = torch.device('cuda:' + dev)\n            device = torch.device(\"cuda\")\n            print(\"Sending model to aforementioned device\")\n            model = model.to(device)\n            print(\n                \"Memory Total : \",\n                round(\n                    torch.cuda.get_device_properties(dev_int).total_memory / 1024 ** 3,\n                    1,\n                ),\n                \"GB, Allocated: \",\n                round(torch.cuda.memory_allocated(dev_int) / 1024 ** 3, 1),\n                \"GB, Cached: \",\n                round(torch.cuda.memory_reserved(dev_int) / 1024 ** 3, 1),\n                \"GB\",\n            )\n\n        print(\n            \"Device - Current: %s Count: %d Name: %s Availability: %s\"\n            % (\n                torch.cuda.current_device(),\n                torch.cuda.device_count(),\n                torch.cuda.get_device_name(device),\n                torch.cuda.is_available(),\n            )\n        )\n\n        if not (optimizer is None):\n            # ensuring optimizer is in correct device - https://github.com/pytorch/pytorch/issues/8741\n            optimizer.load_state_dict(optimizer.state_dict())\n\n    else:\n        dev = -1\n        device = torch.device(\"cpu\")\n        model.cpu()\n        amp = False\n        print(\"Since Device is CPU, Mixed Precision Training is set to False\")\n\n    return model, amp, device\n\n\ndef get_class_imbalance_weights(training_data_loader, parameters):\n    \"\"\"\n    This function calculates the penalty that is used for validation loss in multi-class problems\n\n    Args:\n        training_data_loader (torch.utils.data.DataLoader): The training data loader.\n        parameters (dict): The parameters passed by the user yaml.\n\n    Returns:\n        dict: The penalty weights for different classes under consideration.\n    \"\"\"\n    abs_dict = {}  # absolute counts for each class\n    weights_dict = {}  # average for \"weighted averaging\"\n    penalty_dict = None  # penalty for misclassification\n    # basically, do this for segmentation/classification tasks\n\n    if parameters[\"problem_type\"] != \"regression\":\n        penalty_dict = {}\n        for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n            abs_dict[i] = 0\n            penalty_dict[i] = 0\n\n    penalty_loader = training_data_loader\n\n    # get the weights for use for dice loss\n    total_counter = 0\n\n    # For regression dice penalty need not be taken account\n    # For classification this should be calculated on the basis of predicted labels and mask\n    # iterate through full penalty data\n    for _, (subject) in enumerate(penalty_loader):\n\n        # segmentation needs masks to be one-hot encoded\n        if parameters[\"problem_type\"] == \"segmentation\":\n            # accumulate dice weights for each label\n            mask = subject[\"label\"][torchio.DATA]\n            one_hot_mask = one_hot(mask, parameters[\"model\"][\"class_list\"])\n            for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n                currentNumber = torch.nonzero(\n                    one_hot_mask[:, i, :, :, :], as_tuple=False\n                ).size(0)\n                # class-specific non-zero voxels\n                abs_dict[i] += currentNumber\n                # total number of non-zero voxels to be considered\n                total_counter += currentNumber\n\n        # for classification, the value needs to be used directly\n        elif parameters[\"problem_type\"] == \"classification\":\n            # accumulate weights for each label\n            value_to_predict = subject[\"value_0\"][0]\n            for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n                if value_to_predict == i:\n                    abs_dict[i] += 1\n                    # we only want to increase the counter for those subjects that are defined in the class_list\n                    total_counter += 1\n\n    # Normalize class weights\n    weights_dict = {key: val / total_counter for key, val in abs_dict.items()}\n\n    # get the penalty values - abs_dict contains the overall number for each class in the penalty data\n    for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n        penalty = total_counter  # start with the assumption that all the non-zero voxels (segmentation) or activate labels (classification) make up the penalty\n        for j in range(0, len(parameters[\"model\"][\"class_list\"])):\n            if i != j:  # for differing classes, subtract the current weight\n                penalty -= abs_dict[j]\n\n        # finally, the \"penalty\" variable contains the total number of voxels/activations that are not part of the current class\n        # this is to be used to weight the loss function\n        # adding epsilon to avoid division by zero\n        penalty_dict[i] = (penalty + sys.float_info.epsilon) / (\n            total_counter + sys.float_info.epsilon\n        )\n\n    return penalty_dict, weights_dict\n", "path": "GANDLF/utils/tensor.py"}], "after_files": [{"content": "import os, sys\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torchio\n\n\ndef one_hot(segmask_array, class_list):\n    \"\"\"\n    This function creates a one-hot-encoded mask from the segmentation mask Tensor and specified class list\n\n    Args:\n        segmask_array (torch.Tensor): The segmentation mask Tensor.\n        class_list (list): The list of classes based on which one-hot encoding needs to happen.\n\n    Returns:\n        torch.Tensor: The one-hot encoded torch.Tensor\n    \"\"\"\n    batch_size = segmask_array.shape[0]\n    batch_stack = []\n    for b in range(batch_size):\n        one_hot_stack = []\n        segmask_array_iter = segmask_array[b, 0]\n        bin_mask = segmask_array_iter == 0  # initialize bin_mask\n        # this implementation allows users to combine logical operands\n        for _class in class_list:\n            if isinstance(_class, str):\n                if \"||\" in _class:  # special case\n                    class_split = _class.split(\"||\")\n                    bin_mask = segmask_array_iter == int(class_split[0])\n                    for i in range(1, len(class_split)):\n                        bin_mask = bin_mask | (\n                            segmask_array_iter == int(class_split[i])\n                        )\n                elif \"|\" in _class:  # special case\n                    class_split = _class.split(\"|\")\n                    bin_mask = segmask_array_iter == int(class_split[0])\n                    for i in range(1, len(class_split)):\n                        bin_mask = bin_mask | (\n                            segmask_array_iter == int(class_split[i])\n                        )\n                else:\n                    # assume that it is a simple int\n                    bin_mask = segmask_array_iter == int(_class)\n            else:\n                bin_mask = segmask_array_iter == int(_class)\n                bin_mask = bin_mask.long()\n            one_hot_stack.append(bin_mask)\n        one_hot_stack = torch.stack(one_hot_stack)\n        batch_stack.append(one_hot_stack)\n    batch_stack = torch.stack(batch_stack)\n    return batch_stack\n\n\ndef reverse_one_hot(predmask_array, class_list):\n    \"\"\"\n    This function creates a full segmentation mask Tensor from a one-hot-encoded mask and specified class list\n\n    Args:\n        predmask_array (torch.Tensor): The predicted segmentation mask Tensor.\n        class_list (list): The list of classes based on which one-hot encoding needs to happen.\n\n    Returns:\n        torch.Tensor: The final mask torch.Tensor.\n    \"\"\"\n    if isinstance(predmask_array, torch.Tensor):\n        array_to_consider = predmask_array.cpu().numpy()\n    else:\n        array_to_consider = predmask_array\n    idx_argmax = np.argmax(array_to_consider, axis=0)\n    final_mask = 0\n    special_cases_to_check = [\"||\"]\n    special_case_detected = False\n    max_current = 0\n\n    for _class in class_list:\n        for case in special_cases_to_check:\n            if isinstance(_class, str):\n                if case in _class:  # check if any of the special cases are present\n                    special_case_detected = True\n                    # if present, then split the sub-class\n                    class_split = _class.split(case)\n                    for i in class_split:  # find the max for computation later on\n                        if int(i) > max_current:\n                            max_current = int(i)\n\n    if special_case_detected:\n        start_idx = 0\n        if (class_list[0] == 0) or (class_list[0] == \"0\"):\n            start_idx = 1\n\n        final_mask = np.asarray(predmask_array[start_idx, :, :, :], dtype=int)\n        start_idx += 1\n        for i in range(start_idx, len(class_list)):\n            final_mask += np.asarray(\n                predmask_array[0, :, :, :], dtype=int\n            )  # predmask_array[i,:,:,:].long()\n            # temp_sum = torch.sum(output)\n        # output_2 = (max_current - torch.sum(output)) % max_current\n        # test_2 = 1\n    else:\n        for idx, _class in enumerate(class_list):\n            final_mask = final_mask + (idx_argmax == idx) * _class\n    return final_mask\n\n\ndef send_model_to_device(model, amp, device, optimizer):\n    \"\"\"\n    This function reads the environment variable(s) and send model to correct device\n\n    Args:\n        model (torch.nn.Module): The model that needs to be sent to specified device.\n        amp (bool): Whether automatic mixed precision is to be used.\n        device (str): Device type.\n        optimizer (torch.optim): The optimizer for training.\n\n    Returns:\n        torch.nn.Module: The model after it has been sent to specified device\n        bool: Whether automatic mixed precision is to be used or not.\n        torch.device: Device type.\n    \"\"\"\n    if device != \"cpu\":\n        if os.environ.get(\"CUDA_VISIBLE_DEVICES\") is None:\n            sys.exit(\n                \"Please set the environment variable 'CUDA_VISIBLE_DEVICES' correctly before trying to run GANDLF on GPU\"\n            )\n\n        dev = os.environ.get(\"CUDA_VISIBLE_DEVICES\")\n        # multi-gpu support\n        # ###\n        # # https://discuss.pytorch.org/t/cuda-visible-devices-make-gpu-disappear/21439/17?u=sarthakpati\n        # ###\n        if \",\" in dev:\n            device = torch.device(\"cuda\")\n            model = nn.DataParallel(model, \"[\" + dev + \"]\")\n        else:\n            print(\"Device requested via CUDA_VISIBLE_DEVICES: \", dev)\n            print(\"Total number of CUDA devices: \", torch.cuda.device_count())\n\n            # if only a single visible device, it will be indexed as '0'\n            if torch.cuda.device_count() == 1:\n                dev = \"0\"\n\n            dev_int = int(dev)\n            print(\"Device finally used: \", dev)\n            # device = torch.device('cuda:' + dev)\n            device = torch.device(\"cuda\")\n            print(\"Sending model to aforementioned device\")\n            model = model.to(device)\n            print(\n                \"Memory Total : \",\n                round(\n                    torch.cuda.get_device_properties(dev_int).total_memory / 1024 ** 3,\n                    1,\n                ),\n                \"GB, Allocated: \",\n                round(torch.cuda.memory_allocated(dev_int) / 1024 ** 3, 1),\n                \"GB, Cached: \",\n                round(torch.cuda.memory_reserved(dev_int) / 1024 ** 3, 1),\n                \"GB\",\n            )\n\n        print(\n            \"Device - Current: %s Count: %d Name: %s Availability: %s\"\n            % (\n                torch.cuda.current_device(),\n                torch.cuda.device_count(),\n                torch.cuda.get_device_name(device),\n                torch.cuda.is_available(),\n            )\n        )\n\n        if not (optimizer is None):\n            # ensuring optimizer is in correct device - https://github.com/pytorch/pytorch/issues/8741\n            optimizer.load_state_dict(optimizer.state_dict())\n\n    else:\n        dev = -1\n        device = torch.device(\"cpu\")\n        model.cpu()\n        amp = False\n        print(\"Since Device is CPU, Mixed Precision Training is set to False\")\n\n    return model, amp, device\n\n\ndef get_class_imbalance_weights(training_data_loader, parameters):\n    \"\"\"\n    This function calculates the penalty that is used for validation loss in multi-class problems\n\n    Args:\n        training_data_loader (torch.utils.data.DataLoader): The training data loader.\n        parameters (dict): The parameters passed by the user yaml.\n\n    Returns:\n        dict: The penalty weights for different classes under consideration.\n    \"\"\"\n    abs_dict = {}  # absolute counts for each class\n    weights_dict = {}  # average for \"weighted averaging\"\n    penalty_dict = None  # penalty for misclassification\n    # basically, do this for segmentation/classification tasks\n\n    if parameters[\"problem_type\"] != \"regression\":\n        penalty_dict = {}\n        for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n            abs_dict[i] = 0\n            penalty_dict[i] = 0\n\n    penalty_loader = training_data_loader\n\n    # get the weights for use for dice loss\n    total_counter = 0\n\n    # For regression dice penalty need not be taken account\n    # For classification this should be calculated on the basis of predicted labels and mask\n    # iterate through full penalty data\n    for _, (subject) in enumerate(penalty_loader):\n\n        # segmentation needs masks to be one-hot encoded\n        if parameters[\"problem_type\"] == \"segmentation\":\n            # accumulate dice weights for each label\n            mask = subject[\"label\"][torchio.DATA]\n            one_hot_mask = one_hot(mask, parameters[\"model\"][\"class_list\"])\n            for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n                currentNumber = torch.nonzero(\n                    one_hot_mask[:, i, :, :, :], as_tuple=False\n                ).size(0)\n                # class-specific non-zero voxels\n                abs_dict[i] += currentNumber\n                # total number of non-zero voxels to be considered\n                total_counter += currentNumber\n\n        # for classification, the value needs to be used directly\n        elif parameters[\"problem_type\"] == \"classification\":\n            # accumulate weights for each label\n            value_to_predict = subject[\"value_0\"][0]\n            for i in range(0, len(parameters[\"model\"][\"class_list\"])):\n                if value_to_predict == i:\n                    abs_dict[i] += 1\n                    # we only want to increase the counter for those subjects that are defined in the class_list\n                    total_counter += 1\n\n    # Normalize class weights\n    weights_dict = {\n        key: (val + sys.float_info.epsilon) / total_counter\n        for key, val in abs_dict.items()\n    }\n\n    # get the raw penalty values\n    penalty = {\n        key: total_counter / (len(abs_dict) * (val + sys.float_info.epsilon))\n        for key, val in abs_dict.items()\n    }\n    # normalize penalty to sum of 1\n    penalty_sum = np.fromiter(penalty.values(), dtype=np.float64).sum()\n    penalty_dict = {\n        key: (val + sys.float_info.epsilon) / penalty_sum\n        for key, val in penalty.items()\n    }\n\n    return penalty_dict, weights_dict\n", "path": "GANDLF/utils/tensor.py"}]}