Add BinaryOutput (#1099)

* Add BinaryOutput

* address comments

---------

Co-authored-by: Marc Romeyn <marcromeyn@gmail.com>

merlin/models/torch/__init__.py

@@ -15,6 +15,14 @@
 #
 
 from merlin.models.torch.batch import Batch, Sequence
-from merlin.models.torch.block import Block, ParallelBlock
+from merlin.models.torch.block import Block
+from merlin.models.torch.outputs.base import ModelOutput
+from merlin.models.torch.outputs.classification import BinaryOutput
 
-__all__ = ["Batch", "Block", "ParallelBlock", "Sequence"]
+__all__ = [
+    "Batch",
+    "BinaryOutput",
+    "Block",
+    "ModelOutput",
+    "Sequence",
+]

merlin/models/torch/outputs/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#

merlin/models/torch/outputs/base.py

@@ -0,0 +1,102 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from typing import Optional, Sequence
+
+import torch
+from torch import nn
+from torchmetrics import Metric
+
+from merlin.models.torch.block import Block
+from merlin.schema import ColumnSchema, Schema
+
+
+class ModelOutput(Block):
+    """A base class for prediction tasks.
+
+    Example usage::
+        >>> schema = ColumnSchema(
+        ...    "target",
+        ...    properties={"domain": {"min": 0, "max": 1}},
+        ...    tags=[Tags.CATEGORICAL, Tags.TARGET]
+        ... )
+        >>> model_output = ModelOutput(
+        ...    nn.LazyLinear(1),
+        ...    nn.Sigmoid(),
+        ...    schema=schema
+        ... )
+        >>> input = torch.randn(3, 2)
+        >>> output = model_output(input)
+        >>> print(output)
+        tensor([[0.5529],
+                [0.3562],
+                [0.7478]], grad_fn=<SigmoidBackward0>)
+
+    Parameters
+    ----------
+    schema: Optional[ColumnSchema]
+        The schema defining the column properties.
+    loss: nn.Module
+        The loss function used for training.
+    metrics: Sequence[Metric]
+        The metrics used for evaluation.
+    name: Optional[str]
+        The name of the model output.
+    """
+
+    def __init__(
+        self,
+        *module: nn.Module,
+        schema: Optional[ColumnSchema] = None,
+        loss: Optional[nn.Module] = None,
+        metrics: Sequence[Metric] = (),
+        name: Optional[str] = None,
+    ):
+        """Initializes a ModelOutput object."""
+        super().__init__(*module, name=name)
+
+        self.loss = loss
+        self.metrics = metrics
+        self.output_schema: Schema = Schema()
+
+        if schema:
+            self.setup_schema(schema)
+        self.create_target_buffer()
+
+    def setup_schema(self, schema: Optional[ColumnSchema]):
+        """Set up the schema for the output.
+
+        Parameters
+        ----------
+        schema: ColumnSchema or None
+            The schema defining the column properties.
+        """
+        self.output_schema = Schema([schema])
+
+    def create_target_buffer(self):
+        self.register_buffer("target", torch.zeros(1, dtype=torch.float32))
+
+    def eval(self):
+        """Sets the module in evaluation mode.
+
+        Returns
+        -------
+        nn.Module
+            The module in evaluation mode.
+        """
+        # Reset target
+        self.target = torch.zeros(1, dtype=torch.float32)
+
+        return self.train(False)

merlin/models/torch/outputs/classification.py

@@ -0,0 +1,73 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from typing import Optional, Sequence
+
+from torch import nn
+from torchmetrics import AUROC, Accuracy, Metric, Precision, Recall
+
+import merlin.dtypes as md
+from merlin.models.torch.outputs.base import ModelOutput
+from merlin.schema import ColumnSchema, Schema
+
+
+class BinaryOutput(ModelOutput):
+    """A prediction block for binary classification.
+
+    Parameters
+    ----------
+    schema: Optional[ColumnSchema])
+        The schema defining the column properties. Default is None.
+    loss: nn.Module
+        The loss function used for training. Default is nn.BCEWithLogitsLoss().
+    metrics: Sequence[Metric]
+        The metrics used for evaluation. Default includes Accuracy, AUROC, Precision, and Recall.
+    """
+
+    def __init__(
+        self,
+        schema: Optional[ColumnSchema] = None,
+        loss: nn.Module = nn.BCEWithLogitsLoss(),
+        metrics: Sequence[Metric] = (
+            Accuracy(task="binary"),
+            AUROC(task="binary"),
+            Precision(task="binary"),
+            Recall(task="binary"),
+        ),
+    ):
+        """Initializes a BinaryOutput object."""
+        super().__init__(
+            nn.LazyLinear(1),
+            nn.Sigmoid(),
+            schema=schema,
+            loss=loss,
+            metrics=metrics,
+        )
+
+    def setup_schema(self, target: Optional[ColumnSchema]):
+        """Set up the schema for the output.
+
+        Parameters
+        ----------
+        target: Optional[ColumnSchema]
+            The schema defining the column properties.
+        """
+        _target = target.with_dtype(md.float32)
+        if "domain" not in target.properties:
+            _target = _target.with_properties(
+                {"domain": {"min": 0, "max": 1, "name": _target.name}},
+            )
+
+        self.output_schema = Schema([_target])

tests/unit/torch/outputs/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#

tests/unit/torch/outputs/test_base.py

@@ -0,0 +1,73 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+import torch
+from torch import nn
+from torchmetrics import AUROC, Accuracy
+
+import merlin.models.torch as mm
+from merlin.models.torch.utils import module_utils
+from merlin.schema import ColumnSchema, Schema, Tags
+
+
+class TestModelOutput:
+    def test_init(self):
+        block = mm.Block()
+        loss = nn.BCEWithLogitsLoss()
+        model_output = mm.ModelOutput(block, loss=loss)
+
+        assert isinstance(model_output, mm.ModelOutput)
+        assert model_output.loss is loss
+        assert model_output.metrics == ()
+        assert model_output.output_schema == Schema()
+
+    def test_identity(self):
+        block = mm.Block()
+        loss = nn.BCEWithLogitsLoss()
+        model_output = mm.ModelOutput(block, loss=loss)
+        inputs = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+
+        outputs = module_utils.module_test(model_output, inputs)
+
+        assert torch.equal(inputs, outputs)
+
+    def test_setup_metrics(self):
+        block = mm.Block()
+        loss = nn.BCEWithLogitsLoss()
+        metrics = (Accuracy(task="binary"), AUROC(task="binary"))
+        model_output = mm.ModelOutput(block, loss=loss, metrics=metrics)
+
+        assert model_output.metrics == metrics
+
+    def test_setup_schema(self):
+        block = mm.Block()
+        loss = nn.BCEWithLogitsLoss()
+        schema = ColumnSchema("feature", dtype=np.int32, tags=[Tags.CONTINUOUS])
+        model_output = mm.ModelOutput(block, loss=loss, schema=schema)
+
+        assert isinstance(model_output.output_schema, Schema)
+        assert model_output.output_schema.first == schema
+
+    def test_eval_resets_target(self):
+        block = mm.Block()
+        loss = nn.BCEWithLogitsLoss()
+        model_output = mm.ModelOutput(block, loss=loss)
+
+        assert torch.equal(model_output.target, torch.zeros(1))
+        model_output.target = torch.ones(1)
+        assert torch.equal(model_output.target, torch.ones(1))
+        model_output.eval()
+        assert torch.equal(model_output.target, torch.zeros(1))

tests/unit/torch/outputs/test_classification.py

@@ -0,0 +1,81 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import torch
+from torch import nn
+from torchmetrics import AUROC, Accuracy, Precision, Recall
+from torchmetrics.classification import BinaryF1Score
+
+import merlin.dtypes as md
+import merlin.models.torch as mm
+from merlin.models.torch.utils import module_utils
+from merlin.schema import ColumnSchema, Schema
+
+
+class TestBinaryOutput:
+    def test_init(self):
+        binary_output = mm.BinaryOutput()
+
+        assert isinstance(binary_output, mm.BinaryOutput)
+        assert isinstance(binary_output.loss, nn.BCEWithLogitsLoss)
+        assert binary_output.metrics == (
+            Accuracy(task="binary"),
+            AUROC(task="binary"),
+            Precision(task="binary"),
+            Recall(task="binary"),
+        )
+        assert binary_output.output_schema == Schema()
+
+    def test_identity(self):
+        binary_output = mm.BinaryOutput()
+        inputs = torch.randn(3, 2)
+
+        outputs = module_utils.module_test(binary_output, inputs)
+
+        assert outputs.shape == (3, 1)
+
+    def test_setup_schema(self):
+        schema = ColumnSchema("foo")
+        binary_output = mm.BinaryOutput(schema=schema)
+
+        assert isinstance(binary_output.output_schema, Schema)
+        assert binary_output.output_schema.first.dtype == md.float32
+        assert binary_output.output_schema.first.properties["domain"]["name"] == "foo"
+        assert binary_output.output_schema.first.properties["domain"]["min"] == 0
+        assert binary_output.output_schema.first.properties["domain"]["max"] == 1
+
+    def test_custom_loss(self):
+        binary_output = mm.BinaryOutput(loss=nn.BCELoss())
+        features = torch.randn(3, 2)
+        targets = torch.randint(2, (3, 1), dtype=torch.float32)
+
+        outputs = module_utils.module_test(binary_output, features)
+
+        assert torch.allclose(
+            binary_output.loss(outputs, targets),
+            nn.BCELoss()(outputs, targets),
+        )
+
+    def test_cutom_metrics(self):
+        binary_output = mm.BinaryOutput(metrics=(BinaryF1Score(),))
+        features = torch.randn(3, 2)
+        targets = torch.randint(2, (3, 1), dtype=torch.float32)
+
+        outputs = module_utils.module_test(binary_output, features)
+
+        assert torch.allclose(
+            binary_output.metrics[0](outputs, targets),
+            BinaryF1Score()(outputs, targets),
+        )