CHANGELOG

CHANGELOG.rst

@@ -103,6 +103,55 @@ Added
             module_device={"input_module": -1, f"{task_name}_pred_head": 0},
             scorer=Scorer(metrics=task_metrics[task_name]),
         )
+* `@senwu`_: Add require_prob_for_eval and require_pred_for_eval to optimize score
+  function performance.
+  (`#92 <https://github.com/SenWu/emmental/pull/92>`_)
+
+.. note::
+
+    The current approach during score the model will store probs and preds which might
+    require a lot of memory resources especially for large datasets. The score function
+    is also used in training. To optimize the score function performance, this PR
+    introduces two new arguments in `EmmentalTask`: `require_prob_for_eval` and
+    `require_pred_for_eval` which automatically selects whether `return_probs` or
+    `return_preds`.
+
+    .. code:: python
+
+        task_name = "Task1"
+        EmmentalTask(
+            name=task_name,
+            module_pool=nn.ModuleDict(
+                {
+                    "input_module": nn.Linear(2, 8),
+                    f"{task_name}_pred_head": nn.Linear(8, 2),
+                }
+            ),
+            task_flow=[
+                {
+                    "name": "input",
+                    "module": "input_module",
+                    "inputs": [("_input_", "data")],
+                },
+                {
+                    "name": f"{task_name}_pred_head",
+                    "module": f"{task_name}_pred_head",
+                    "inputs": [("input", 0)],
+                },
+            ],
+            loss_func=partial(ce_loss, task_name),
+            output_func=partial(output, task_name),
+            action_outputs=[
+                (f"{task_name}_pred_head", 0),
+                ("_input_", "data"),
+                (f"{task_name}_pred_head", 0),
+            ],
+            module_device={"input_module": -1, f"{task_name}_pred_head": 0},
+            require_prob_for_eval=True,
+            require_pred_for_eval=True,
+            scorer=Scorer(metrics=task_metrics[task_name]),
+        )
+
 
 Fixed
 ^^^^^