Add DLRM Model (#1171)

* Add DLRM Model

* make model a class rather a function

---------

Co-authored-by: Marc Romeyn <[email protected]>

merlin/models/torch/__init__.py

@@ -23,6 +23,7 @@
 from merlin.models.torch.inputs.select import SelectFeatures, SelectKeys
 from merlin.models.torch.inputs.tabular import TabularInputBlock
 from merlin.models.torch.models.base import Model
+from merlin.models.torch.models.ranking import DLRMModel
 from merlin.models.torch.outputs.base import ModelOutput
 from merlin.models.torch.outputs.classification import BinaryOutput
 from merlin.models.torch.outputs.regression import RegressionOutput
@@ -55,4 +56,5 @@
     "Stack",
     "schema",
     "DLRMBlock",
+    "DLRMModel",
 ]

merlin/models/torch/models/base.py

@@ -196,6 +196,11 @@ def compute_loss(
         else:
             raise ValueError(f"Unknown 'predictions' type: {type(predictions)}")
 
+        if _targets.size() != _predictions.size():
+            _targets = _targets.view(_predictions.size())
+        if _targets.type() != _predictions.type():
+            _targets = _targets.type_as(_predictions)
+
         results["loss"] = results["loss"] + model_out.loss(_predictions, _targets) / len(
             model_outputs
         )

merlin/models/torch/models/ranking.py

@@ -0,0 +1,76 @@
+from typing import Optional
+
+from torch import nn
+
+from merlin.models.torch.block import Block
+from merlin.models.torch.blocks.dlrm import DLRMBlock
+from merlin.models.torch.models.base import Model
+from merlin.models.torch.outputs.tabular import TabularOutputBlock
+from merlin.schema import Schema
+
+
+class DLRMModel(Model):
+    """
+    The Deep Learning Recommendation Model (DLRM) as proposed in Naumov, et al. [1]
+
+    Parameters
+    ----------
+    schema : Schema
+        The schema to use for selection.
+    dim : int
+        The dimensionality of the output vectors.
+    bottom_block : Block
+        Block to pass the continuous features to.
+        Note that, the output dimensionality of this block must be equal to ``dim``.
+    top_block : Block, optional
+        An optional upper-level block of the model.
+    interaction : nn.Module, optional
+        Interaction module for DLRM.
+        If not provided, DLRMInteraction will be used by default.
+    output_block : Block, optional
+        The output block of the model, by default None.
+        If None, a TabularOutputBlock with schema and default initializations is used.
+
+    Returns
+    -------
+    Model
+        An instance of Model class representing the fully formed DLRM.
+
+    Example usage
+    -------------
+    >>> model = mm.DLRMModel(
+    ...    schema,
+    ...    dim=64,
+    ...    bottom_block=mm.MLPBlock([256, 64]),
+    ...    output_block=BinaryOutput(ColumnSchema("target")))
+    >>> trainer = pl.Trainer()
+    >>> model.initialize(dataloader)
+    >>> trainer.fit(model, dataloader)
+
+    References
+    ----------
+    [1] Naumov, Maxim, et al. "Deep learning recommendation model for
+        personalization and recommendation systems." arXiv preprint arXiv:1906.00091 (2019).
+    """
+
+    def __init__(
+        self,
+        schema: Schema,
+        dim: int,
+        bottom_block: Block,
+        top_block: Optional[Block] = None,
+        interaction: Optional[nn.Module] = None,
+        output_block: Optional[Block] = None,
+    ) -> None:
+        if output_block is None:
+            output_block = TabularOutputBlock(schema, init="defaults")
+
+        dlrm_body = DLRMBlock(
+            schema,
+            dim,
+            bottom_block,
+            top_block=top_block,
+            interaction=interaction,
+        )
+
+        super().__init__(dlrm_body, output_block)

merlin/models/torch/utils/module_utils.py

@@ -236,7 +236,7 @@ def initialize(module, data: Union[Dataset, Loader, Batch], dtype=torch.float32)
     if hasattr(module, "model_outputs"):
         for model_out in module.model_outputs():
             for metric in model_out.metrics:
-                metric.to(batch.device())
+                metric.to(device=batch.device())
 
     from merlin.models.torch import schema
 

tests/unit/torch/models/test_base.py

@@ -15,14 +15,15 @@
 #
 import pandas as pd
 import pytest
+import pytorch_lightning as pl
 import torch
 from torch import nn
 from torchmetrics import AUROC, Accuracy, Precision, Recall
 
 import merlin.models.torch as mm
 from merlin.dataloader.torch import Loader
 from merlin.io import Dataset
-from merlin.models.torch.batch import Batch
+from merlin.models.torch.batch import Batch, sample_batch
 from merlin.models.torch.models.base import compute_loss
 from merlin.models.torch.utils import module_utils
 from merlin.schema import ColumnSchema
@@ -200,22 +201,29 @@ def test_no_output_schema(self):
         with pytest.raises(ValueError, match="Could not get output schema of PlusOne()"):
             mm.schema.output(model)
 
-    # def test_train_classification(self, music_streaming_data):
-    #     schema = music_streaming_data.schema.without(["user_genres", "like", "item_genres"])
-    #     music_streaming_data.schema = schema
+    def test_train_classification_with_lightning_trainer(self, music_streaming_data, batch_size=16):
+        schema = music_streaming_data.schema.select_by_name(
+            ["item_id", "user_id", "user_age", "item_genres", "click"]
+        )
+        music_streaming_data.schema = schema
 
-    #     model = mm.Model(
-    #        mm.TabularInputBlock(schema),
-    #        mm.MLPBlock([4, 2]),
-    #        mm.BinaryOutput(schema.select_by_name("click").first),
-    #        schema=schema,
-    #     )
+        model = mm.Model(
+            mm.TabularInputBlock(schema, init="defaults"),
+            mm.MLPBlock([4, 2]),
+            mm.BinaryOutput(schema.select_by_name("click").first),
+        )
+
+        trainer = pl.Trainer(max_epochs=1, devices=1)
+
+        with Loader(music_streaming_data, batch_size=batch_size) as loader:
+            model.initialize(loader)
+            trainer.fit(model, loader)
 
-    #     trainer = pl.Trainer(max_epochs=1)
+        assert trainer.logged_metrics["train_loss"] > 0.0
+        assert trainer.num_training_batches == 7  # 100 rows // 16 per batch + 1 for last batch
 
-    #     with Loader(music_streaming_data, batch_size=16) as loader:
-    #         model.initialize(loader)
-    #         trainer.fit(model, loader)
+        batch = sample_batch(music_streaming_data, batch_size)
+        _ = module_utils.module_test(model, batch)
 
 
 class TestComputeLoss:

tests/unit/torch/models/test_ranking.py

@@ -0,0 +1,38 @@
+import pytest
+import pytorch_lightning as pl
+
+import merlin.models.torch as mm
+from merlin.dataloader.torch import Loader
+from merlin.models.torch.batch import sample_batch
+from merlin.models.torch.utils import module_utils
+from merlin.schema import ColumnSchema
+
+
+@pytest.mark.parametrize("output_block", [None, mm.BinaryOutput(ColumnSchema("click"))])
+class TestDLRMModel:
+    def test_train_dlrm_with_lightning_loader(
+        self, music_streaming_data, output_block, dim=2, batch_size=16
+    ):
+        schema = music_streaming_data.schema.select_by_name(
+            ["item_id", "user_id", "user_age", "item_genres", "click"]
+        )
+        music_streaming_data.schema = schema
+
+        model = mm.DLRMModel(
+            schema,
+            dim=dim,
+            bottom_block=mm.MLPBlock([4, 2]),
+            top_block=mm.MLPBlock([4, 2]),
+            output_block=output_block,
+        )
+
+        trainer = pl.Trainer(max_epochs=1, devices=1)
+
+        with Loader(music_streaming_data, batch_size=batch_size) as train_loader:
+            model.initialize(train_loader)
+            trainer.fit(model, train_loader)
+
+        assert trainer.logged_metrics["train_loss"] > 0.0
+
+        batch = sample_batch(music_streaming_data, batch_size)
+        _ = module_utils.module_test(model, batch)