device fix and docs

Author: blondered

rectools/models/nn/transformers/base.py

@@ -176,7 +176,7 @@ class TransformerModelConfig(ModelConfig):
     verbose: int = 0
     deterministic: bool = False
     recommend_batch_size: int = 256
-    recommend_device: tp.Optional[str] = None
+    recommend_torch_device: tp.Optional[str] = None
     recommend_n_threads: int = 0
     recommend_use_torch_ranking: bool = True
     train_min_user_interactions: int = 2
@@ -233,7 +233,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         verbose: int = 0,
         deterministic: bool = False,
         recommend_batch_size: int = 256,
-        recommend_device: tp.Optional[str] = None,
+        recommend_torch_device: tp.Optional[str] = None,
         recommend_n_threads: int = 0,
         recommend_use_torch_ranking: bool = True,
         train_min_user_interactions: int = 2,
@@ -270,7 +270,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         self.epochs = epochs
         self.deterministic = deterministic
         self.recommend_batch_size = recommend_batch_size
-        self.recommend_device = recommend_device
+        self.recommend_torch_device = recommend_torch_device
         self.recommend_n_threads = recommend_n_threads
         self.recommend_use_torch_ranking = recommend_use_torch_ranking
         self.train_min_user_interactions = train_min_user_interactions
@@ -454,7 +454,7 @@ def _recommend_u2i(
             dataset=dataset,
             n_threads=self.recommend_n_threads,
             use_torch_ranking=self.recommend_use_torch_ranking,
-            device=self.recommend_device,
+            torch_device=self.recommend_torch_device,
         )
 
     def _recommend_i2i(
@@ -473,7 +473,7 @@ def _recommend_i2i(
             k=k,
             n_threads=self.recommend_n_threads,
             use_torch_ranking=self.recommend_use_torch_ranking,
-            device=self.recommend_device,
+            torch_device=self.recommend_torch_device,
         )
 
     @property

rectools/models/nn/transformers/bert4rec.py

@@ -269,19 +269,24 @@ class BERT4RecModel(TransformerModelBase[BERT4RecModelConfig]):
         How many samples per batch to load during `recommend`.
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_batch_size` attribute.
-    recommend_device : {"cpu", "cuda", "cuda:0", ...}, default ``None``
-        String representation for `torch.device` used for recommendations.
+    recommend_torch_device : {"cpu", "cuda", "cuda:0", ...}, default ``None``
+        String representation for `torch.device` used for torch model inference.
+        When `recommend_use_torch_ranking` is set to ``True`` (default) this device is also used
+        for items ranking while preparing recommendations using `TorchRanker`.
         When set to ``None``, "cuda" will be used if it is available, "cpu" otherwise.
         If you want to change this parameter after model is initialized,
-        you can manually assign new value to model `recommend_device` attribute.
+        you can manually assign new value to model `recommend_torch_device` attribute.
     recommend_use_torch_ranking : bool, default ``True``
         Use `TorchRanker` for items ranking while preparing recommendations.
-        If set to ``False``, use `ImplicitRanker` instead.
+        When set to ``True`` (default), device specified in `recommend_torch_device` is used
+        for items ranking.
+        When set to ``False``, multi-threaded cpu ranking will be used with `ImplicitRanker`. You
+        can specify numer of threads using `recommend_n_threads` argument.
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_use_torch_ranking` attribute.
     recommend_n_threads : int, default 0
-        Number of threads to use for `ImplicitRanker`. Omitted if `recommend_use_torch_ranking` is
-        set to ``True`` (default).
+        Number of threads to use for cpu items ranking with `ImplicitRanker`. Omitted if
+        `recommend_use_torch_ranking` is set to ``True`` (default).
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_n_threads` attribute.
     data_preparator_kwargs: optional(dict), default ``None``
@@ -333,7 +338,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         get_val_mask_func: tp.Optional[ValMaskCallable] = None,
         get_trainer_func: tp.Optional[TrainerCallable] = None,
         recommend_batch_size: int = 256,
-        recommend_device: tp.Optional[str] = None,
+        recommend_torch_device: tp.Optional[str] = None,
         recommend_use_torch_ranking: bool = True,
         recommend_n_threads: int = 0,
         data_preparator_kwargs: tp.Optional[InitKwargs] = None,
@@ -366,7 +371,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
             verbose=verbose,
             deterministic=deterministic,
             recommend_batch_size=recommend_batch_size,
-            recommend_device=recommend_device,
+            recommend_torch_device=recommend_torch_device,
             recommend_n_threads=recommend_n_threads,
             recommend_use_torch_ranking=recommend_use_torch_ranking,
             train_min_user_interactions=train_min_user_interactions,

rectools/models/nn/transformers/lightning.py

@@ -115,7 +115,7 @@ def _recommend_u2i(
         filter_viewed: bool,
         use_torch_ranking: bool,
         n_threads: int,
-        device: tp.Optional[str],
+        torch_device: tp.Optional[str],
         *args: tp.Any,
         **kwargs: tp.Any,
     ) -> InternalRecoTriplet:
@@ -129,7 +129,7 @@ def _recommend_i2i(
         k: int,
         use_torch_ranking: bool,
         n_threads: int,
-        device: tp.Optional[str],
+        torch_device: tp.Optional[str],
         *args: tp.Any,
         **kwargs: tp.Any,
     ) -> InternalRecoTriplet:
@@ -295,23 +295,23 @@ def _xavier_normal_init(self) -> None:
             if param.data.dim() > 1:
                 torch.nn.init.xavier_normal_(param.data)
 
-    def _prepare_for_inference(self, recommend_device: tp.Optional[str]) -> None:
-        if recommend_device is None:
-            recommend_device = "cuda" if torch.cuda.is_available() else "cpu"
-        device = torch.device(recommend_device)
+    def _prepare_for_inference(self, torch_device: tp.Optional[str]) -> None:
+        if torch_device is None:
+            torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = torch.device(torch_device)
         self.torch_model.to(device)
         self.torch_model.eval()
 
     def _get_user_item_embeddings(
         self,
         recommend_dataloader: DataLoader,
-        recommend_device: tp.Optional[str],
+        torch_device: tp.Optional[str],
     ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
         """
         Prepare user embeddings for all user interaction sequences in `recommend_dataloader`.
         Prepare item embeddings for full items catalog.
         """
-        self._prepare_for_inference(recommend_device)
+        self._prepare_for_inference(torch_device)
         device = self.torch_model.item_model.device
 
         with torch.no_grad():
@@ -333,14 +333,14 @@ def _recommend_u2i(
         filter_viewed: bool,
         use_torch_ranking: bool,
         n_threads: int,
-        device: tp.Optional[str],
+        torch_device: tp.Optional[str],
     ) -> InternalRecoTriplet:
         """Recommend to users."""
         ui_csr_for_filter = None
         if filter_viewed:
             ui_csr_for_filter = dataset.get_user_item_matrix(include_weights=False, include_warm_items=True)[user_ids]
 
-        user_embs, item_embs = self._get_user_item_embeddings(recommend_dataloader, device)
+        user_embs, item_embs = self._get_user_item_embeddings(recommend_dataloader, torch_device)
 
         ranker: Ranker
         if use_torch_ranking:
@@ -379,10 +379,10 @@ def _recommend_i2i(
         k: int,
         use_torch_ranking: bool,
         n_threads: int,
-        device: tp.Optional[str],
+        torch_device: tp.Optional[str],
     ) -> InternalRecoTriplet:
         """Recommend to items."""
-        self._prepare_for_inference(device)
+        self._prepare_for_inference(torch_device)
         with torch.no_grad():
             item_embs = self.torch_model.item_model.get_all_embeddings()
 

rectools/models/nn/transformers/sasrec.py

@@ -349,19 +349,24 @@ class SASRecModel(TransformerModelBase[SASRecModelConfig]):
         How many samples per batch to load during `recommend`.
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_batch_size` attribute.
-    recommend_device : {"cpu", "cuda", "cuda:0", ...}, default ``None``
-        String representation for `torch.device` used for recommendations.
+    recommend_torch_device : {"cpu", "cuda", "cuda:0", ...}, default ``None``
+        String representation for `torch.device` used for torch model inference.
+        When `recommend_use_torch_ranking` is set to ``True`` (default) this device is also used
+        for items ranking while preparing recommendations using `TorchRanker`.
         When set to ``None``, "cuda" will be used if it is available, "cpu" otherwise.
         If you want to change this parameter after model is initialized,
-        you can manually assign new value to model `recommend_device` attribute.
+        you can manually assign new value to model `recommend_torch_device` attribute.
     recommend_use_torch_ranking : bool, default ``True``
         Use `TorchRanker` for items ranking while preparing recommendations.
-        If set to ``False``, use `ImplicitRanker` instead.
+        When set to ``True`` (default), device specified in `recommend_torch_device` is used
+        for items ranking.
+        When set to ``False``, multi-threaded cpu ranking will be used with `ImplicitRanker`. You
+        can specify numer of threads using `recommend_n_threads` argument.
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_use_torch_ranking` attribute.
     recommend_n_threads : int, default 0
-        Number of threads to use for `ImplicitRanker`. Omitted if `recommend_use_torch_ranking` is
-        set to ``True`` (default).
+        Number of threads to use for cpu items ranking with `ImplicitRanker`. Omitted if
+        `recommend_use_torch_ranking` is set to ``True`` (default).
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_n_threads` attribute.
     data_preparator_kwargs: optional(dict), default ``None``
@@ -412,7 +417,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         get_val_mask_func: tp.Optional[ValMaskCallable] = None,
         get_trainer_func: tp.Optional[TrainerCallable] = None,
         recommend_batch_size: int = 256,
-        recommend_device: tp.Optional[str] = None,
+        recommend_torch_device: tp.Optional[str] = None,
         recommend_use_torch_ranking: bool = True,
         recommend_n_threads: int = 0,
         data_preparator_kwargs: tp.Optional[InitKwargs] = None,
@@ -442,7 +447,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
             verbose=verbose,
             deterministic=deterministic,
             recommend_batch_size=recommend_batch_size,
-            recommend_device=recommend_device,
+            recommend_torch_device=recommend_torch_device,
             recommend_n_threads=recommend_n_threads,
             recommend_use_torch_ranking=recommend_use_torch_ranking,
             train_min_user_interactions=train_min_user_interactions,

tests/models/nn/transformers/test_bert4rec.py

@@ -113,7 +113,7 @@ def get_trainer() -> Trainer:
 
     @pytest.mark.parametrize("recommend_use_torch_ranking", (True, False))
     @pytest.mark.parametrize(
-        "accelerator,n_devices,recommend_device",
+        "accelerator,n_devices,recommend_torch_device",
         [
             ("cpu", 1, "cpu"),
             pytest.param(
@@ -219,7 +219,7 @@ def test_u2i(
         filter_viewed: bool,
         accelerator: str,
         n_devices: int,
-        recommend_device: str,
+        recommend_torch_device: str,
         expected_cpu_1: pd.DataFrame,
         expected_cpu_2: pd.DataFrame,
         expected_gpu_1: pd.DataFrame,
@@ -248,7 +248,7 @@ def get_trainer() -> Trainer:
             batch_size=4,
             epochs=2,
             deterministic=True,
-            recommend_device=recommend_device,
+            recommend_torch_device=recommend_torch_device,
             item_net_block_types=(IdEmbeddingsItemNet,),
             get_trainer_func=get_trainer,
             recommend_use_torch_ranking=recommend_use_torch_ranking,
@@ -832,7 +832,7 @@ def initial_config(self) -> tp.Dict[str, tp.Any]:
             "epochs": 10,
             "verbose": 1,
             "deterministic": True,
-            "recommend_device": None,
+            "recommend_torch_device": None,
             "recommend_batch_size": 256,
             "recommend_n_threads": 0,
             "recommend_use_torch_ranking": True,

tests/models/nn/transformers/test_sasrec.py

@@ -158,7 +158,7 @@ def get_trainer() -> Trainer:
 
     @pytest.mark.parametrize("recommend_use_torch_ranking", (True, False))
     @pytest.mark.parametrize(
-        "accelerator,devices,recommend_device",
+        "accelerator,devices,recommend_torch_device",
         [
             ("cpu", 1, "cpu"),
             pytest.param(
@@ -250,7 +250,7 @@ def test_u2i(
         filter_viewed: bool,
         accelerator: str,
         devices: tp.Union[int, tp.List[int]],
-        recommend_device: str,
+        recommend_torch_device: str,
         expected_cpu_1: pd.DataFrame,
         expected_cpu_2: pd.DataFrame,
         expected_gpu: pd.DataFrame,
@@ -279,7 +279,7 @@ def get_trainer() -> Trainer:
             batch_size=4,
             epochs=2,
             deterministic=True,
-            recommend_device=recommend_device,
+            recommend_torch_device=recommend_torch_device,
             item_net_block_types=(IdEmbeddingsItemNet,),
             get_trainer_func=get_trainer,
             recommend_use_torch_ranking=recommend_use_torch_ranking,
@@ -910,7 +910,7 @@ def initial_config(self) -> tp.Dict[str, tp.Any]:
             "epochs": 10,
             "verbose": 1,
             "deterministic": True,
-            "recommend_device": None,
+            "recommend_torch_device": None,
             "recommend_batch_size": 256,
             "recommend_n_threads": 0,
             "recommend_use_torch_ranking": True,