autogluon.tabular 1.3.2b20250610__py3-none-any.whl → 1.4.1b20251214__py3-none-any.whl

autogluon/tabular/models/tabm/tabm_reference.py

@@ -0,0 +1,631 @@
+# License: https://github.com/yandex-research/tabm/blob/main/LICENSE
+
+# NOTE
+# The minimum required versions of the dependencies are specified in README.md.
+
+from __future__ import annotations
+
+import itertools
+from typing import Any, Literal, Union
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from . import rtdl_num_embeddings
+from .rtdl_num_embeddings import _Periodic
+
+
+# ======================================================================================
+# Initialization
+# ======================================================================================
+def init_rsqrt_uniform_(x: Tensor, d: int) -> Tensor:
+    assert d > 0
+    d_rsqrt = d**-0.5
+    return nn.init.uniform_(x, -d_rsqrt, d_rsqrt)
+
+
+@torch.inference_mode()
+def init_random_signs_(x: Tensor) -> Tensor:
+    return x.bernoulli_(0.5).mul_(2).add_(-1)
+
+
+# ======================================================================================
+# Modules
+# ======================================================================================
+class NLinear(nn.Module):
+    """N linear layers applied in parallel to N disjoint parts of the input.
+
+    **Shape**
+
+    - Input: ``(B, N, in_features)``
+    - Output: ``(B, N, out_features)``
+
+    The i-th linear layer is applied to the i-th matrix of the shape (B, in_features).
+
+    Technically, this is a simplified version of delu.nn.NLinear:
+    https://yura52.github.io/delu/stable/api/generated/delu.nn.NLinear.html.
+    The difference is that this layer supports only 3D inputs
+    with exactly one batch dimension. By contrast, delu.nn.NLinear supports
+    any number of batch dimensions.
+    """
+
+    def __init__(
+        self, n: int, in_features: int, out_features: int, bias: bool = True
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(n, in_features, out_features))
+        self.bias = nn.Parameter(torch.empty(n, out_features)) if bias else None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        d = self.weight.shape[-2]
+        init_rsqrt_uniform_(self.weight, d)
+        if self.bias is not None:
+            init_rsqrt_uniform_(self.bias, d)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.ndim == 3
+        assert x.shape[-(self.weight.ndim - 1) :] == self.weight.shape[:-1]
+
+        x = x.transpose(0, 1)
+        x = x @ self.weight
+        x = x.transpose(0, 1)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+class OneHotEncoding0d(nn.Module):
+    # Input:  (*, n_cat_features=len(cardinalities))
+    # Output: (*, sum(cardinalities))
+
+    def __init__(self, cardinalities: list[int]) -> None:
+        super().__init__()
+        self._cardinalities = cardinalities
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert x.ndim >= 1
+        assert x.shape[-1] == len(self._cardinalities)
+
+        return torch.cat(
+            [
+                # NOTE
+                # This is a quick hack to support out-of-vocabulary categories.
+                #
+                # Recall that lib.data.transform_cat encodes categorical features
+                # as follows:
+                # - In-vocabulary values receive indices from `range(cardinality)`.
+                # - All out-of-vocabulary values (i.e. new categories in validation
+                #   and test data that are not presented in the training data)
+                #   receive the index `cardinality`.
+                #
+                # As such, the line below will produce the standard one-hot encoding for
+                # known categories, and the all-zeros encoding for unknown categories.
+                # This may not be the best approach to deal with unknown values,
+                # but should be enough for our purposes.
+                nn.functional.one_hot(x[..., i], cardinality + 1)[..., :-1]
+                for i, cardinality in enumerate(self._cardinalities)
+            ],
+            -1,
+        )
+
+
+class ScaleEnsemble(nn.Module):
+    def __init__(
+        self,
+        k: int,
+        d: int,
+        *,
+        init: Literal['ones', 'normal', 'random-signs'],
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(k, d))
+        self._weight_init = init
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self._weight_init == 'ones':
+            nn.init.ones_(self.weight)
+        elif self._weight_init == 'normal':
+            nn.init.normal_(self.weight)
+        elif self._weight_init == 'random-signs':
+            init_random_signs_(self.weight)
+        else:
+            raise ValueError(f'Unknown weight_init: {self._weight_init}')
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert x.ndim >= 2
+        return x * self.weight
+
+
+class LinearEfficientEnsemble(nn.Module):
+    """
+    This layer is a more configurable version of the "BatchEnsemble" layer
+    from the paper
+    "BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning"
+    (link: https://arxiv.org/abs/2002.06715).
+
+    First, this layer allows to select only some of the "ensembled" parts:
+    - the input scaling  (r_i in the BatchEnsemble paper)
+    - the output scaling (s_i in the BatchEnsemble paper)
+    - the output bias    (not mentioned in the BatchEnsemble paper,
+                          but is presented in public implementations)
+
+    Second, the initialization of the scaling weights is configurable
+    through the `scaling_init` argument.
+
+    NOTE
+    The term "adapter" is used in the TabM paper only to tell the story.
+    The original BatchEnsemble paper does NOT use this term. So this class also
+    avoids the term "adapter".
+    """
+
+    r: Union[None, Tensor]
+    s: Union[None, Tensor]
+    bias: Union[None, Tensor]
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        *,
+        k: int,
+        ensemble_scaling_in: bool,
+        ensemble_scaling_out: bool,
+        ensemble_bias: bool,
+        scaling_init: Literal['ones', 'random-signs'],
+    ):
+        assert k > 0
+        if ensemble_bias:
+            assert bias
+        super().__init__()
+
+        self.weight = nn.Parameter(torch.empty(out_features, in_features))
+        self.register_parameter(
+            'r',
+            (
+                nn.Parameter(torch.empty(k, in_features))
+                if ensemble_scaling_in
+                else None
+            ),  # type: ignore[code]
+        )
+        self.register_parameter(
+            's',
+            (
+                nn.Parameter(torch.empty(k, out_features))
+                if ensemble_scaling_out
+                else None
+            ),  # type: ignore[code]
+        )
+        self.register_parameter(
+            'bias',
+            (
+                nn.Parameter(torch.empty(out_features))  # type: ignore[code]
+                if bias and not ensemble_bias
+                else nn.Parameter(torch.empty(k, out_features))
+                if ensemble_bias
+                else None
+            ),
+        )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.k = k
+        self.scaling_init = scaling_init
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init_rsqrt_uniform_(self.weight, self.in_features)
+        scaling_init_fn = {'ones': nn.init.ones_, 'random-signs': init_random_signs_}[
+            self.scaling_init
+        ]
+        if self.r is not None:
+            scaling_init_fn(self.r)
+        if self.s is not None:
+            scaling_init_fn(self.s)
+        if self.bias is not None:
+            bias_init = torch.empty(
+                # NOTE: the shape of bias_init is (out_features,) not (k, out_features).
+                # It means that all biases have the same initialization.
+                # This is similar to having one shared bias plus
+                # k zero-initialized non-shared biases.
+                self.out_features,
+                dtype=self.weight.dtype,
+                device=self.weight.device,
+            )
+            bias_init = init_rsqrt_uniform_(bias_init, self.in_features)
+            with torch.inference_mode():
+                self.bias.copy_(bias_init)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # x.shape == (B, K, D)
+        assert x.ndim == 3
+
+        # >>> The equation (5) from the BatchEnsemble paper (arXiv v2).
+        if self.r is not None:
+            x = x * self.r
+        x = x @ self.weight.T
+        if self.s is not None:
+            x = x * self.s
+        # <<<
+
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        *,
+        d_in: Union[None, int] = None,
+        d_out: Union[None, int] = None,
+        n_blocks: int,
+        d_block: int,
+        dropout: float,
+        activation: str = 'ReLU',
+    ) -> None:
+        super().__init__()
+
+        d_first = d_block if d_in is None else d_in
+        self.blocks = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(d_first if i == 0 else d_block, d_block),
+                    getattr(nn, activation)(),
+                    nn.Dropout(dropout),
+                )
+                for i in range(n_blocks)
+            ]
+        )
+        self.output = None if d_out is None else nn.Linear(d_block, d_out)
+
+    def forward(self, x: Tensor) -> Tensor:
+        for block in self.blocks:
+            x = block(x)
+        if self.output is not None:
+            x = self.output(x)
+        return x
+
+
+def make_efficient_ensemble(module: nn.Module, EnsembleLayer, **kwargs) -> None:
+    """Replace linear layers with efficient ensembles of linear layers.
+
+    NOTE
+    In the paper, there are no experiments with networks with normalization layers.
+    Perhaps, their trainable weights (the affine transformations) also need
+    "ensemblification" as in the paper about "FiLM-Ensemble".
+    Additional experiments are required to make conclusions.
+    """
+    for name, submodule in list(module.named_children()):
+        if isinstance(submodule, nn.Linear):
+            module.add_module(
+                name,
+                EnsembleLayer(
+                    in_features=submodule.in_features,
+                    out_features=submodule.out_features,
+                    bias=submodule.bias is not None,
+                    **kwargs,
+                ),
+            )
+        else:
+            make_efficient_ensemble(submodule, EnsembleLayer, **kwargs)
+
+
+def _get_first_ensemble_layer(backbone: MLP) -> LinearEfficientEnsemble:
+    if isinstance(backbone, MLP):
+        return backbone.blocks[0][0]  # type: ignore[code]
+    else:
+        raise RuntimeError(f'Unsupported backbone: {backbone}')
+
+
+@torch.inference_mode()
+def _init_first_adapter(
+    weight: Tensor,
+    distribution: Literal['normal', 'random-signs'],
+    init_sections: list[int],
+) -> None:
+    """Initialize the first adapter.
+
+    NOTE
+    The `init_sections` argument is a historical artifact that accidentally leaked
+    from irrelevant experiments to the final models. Perhaps, the code related
+    to `init_sections` can be simply removed, but this was not tested.
+    """
+    assert weight.ndim == 2
+    assert weight.shape[1] == sum(init_sections)
+
+    if distribution == 'normal':
+        init_fn_ = nn.init.normal_
+    elif distribution == 'random-signs':
+        init_fn_ = init_random_signs_
+    else:
+        raise ValueError(f'Unknown distribution: {distribution}')
+
+    section_bounds = [0, *torch.tensor(init_sections).cumsum(0).tolist()]
+    for i in range(len(init_sections)):
+        # NOTE
+        # As noted above, this section-based initialization is an arbitrary historical
+        # artifact. Consider the first adapter of one ensemble member.
+        # This adapter vector is implicitly split into "sections",
+        # where one section corresponds to one feature. The code below ensures that
+        # the adapter weights in one section are initialized with the same random value
+        # from the given distribution.
+        w = torch.empty((len(weight), 1), dtype=weight.dtype, device=weight.device)
+        init_fn_(w)
+        weight[:, section_bounds[i] : section_bounds[i + 1]] = w
+
+
+_CUSTOM_MODULES = {
+    # https://docs.python.org/3/library/stdtypes.html#definition.__name__
+    CustomModule.__name__: CustomModule
+    for CustomModule in [
+        rtdl_num_embeddings.LinearEmbeddings,
+        rtdl_num_embeddings.LinearReLUEmbeddings,
+        rtdl_num_embeddings.PeriodicEmbeddings,
+        rtdl_num_embeddings.PiecewiseLinearEmbeddings,
+        MLP,
+    ]
+}
+
+
+def make_module(type: str, *args, **kwargs) -> nn.Module:
+    Module = getattr(nn, type, None)
+    if Module is None:
+        Module = _CUSTOM_MODULES[type]
+    return Module(*args, **kwargs)
+
+
+# ======================================================================================
+# Optimization
+# ======================================================================================
+def default_zero_weight_decay_condition(
+    module_name: str, module: nn.Module, parameter_name: str, parameter: nn.Parameter
+):
+    del module_name, parameter
+    return parameter_name.endswith('bias') or isinstance(
+        module,
+        (
+            nn.BatchNorm1d,
+            nn.LayerNorm,
+            nn.InstanceNorm1d,
+            rtdl_num_embeddings.LinearEmbeddings,
+            rtdl_num_embeddings.LinearReLUEmbeddings,
+            _Periodic,
+        ),
+    )
+
+
+def make_parameter_groups(
+    module: nn.Module,
+    zero_weight_decay_condition=default_zero_weight_decay_condition,
+    custom_groups: Union[None, list[dict[str, Any]]] = None,
+) -> list[dict[str, Any]]:
+    if custom_groups is None:
+        custom_groups = []
+    custom_params = frozenset(
+        itertools.chain.from_iterable(group['params'] for group in custom_groups)
+    )
+    assert len(custom_params) == sum(
+        len(group['params']) for group in custom_groups
+    ), 'Parameters in custom_groups must not intersect'
+    zero_wd_params = frozenset(
+        p
+        for mn, m in module.named_modules()
+        for pn, p in m.named_parameters()
+        if p not in custom_params and zero_weight_decay_condition(mn, m, pn, p)
+    )
+    default_group = {
+        'params': [
+            p
+            for p in module.parameters()
+            if p not in custom_params and p not in zero_wd_params
+        ]
+    }
+    return [
+        default_group,
+        {'params': list(zero_wd_params), 'weight_decay': 0.0},
+        *custom_groups,
+    ]
+
+
+# ======================================================================================
+# The model
+# ======================================================================================
+class Model(nn.Module):
+    """MLP & TabM."""
+
+    def __init__(
+        self,
+        *,
+        n_num_features: int,
+        cat_cardinalities: list[int],
+        n_classes: Union[None, int],
+        backbone: dict,
+        bins: Union[None, list[Tensor]],  # For piecewise-linear encoding/embeddings.
+        num_embeddings: Union[None, dict] = None,
+        arch_type: Literal[
+            # Plain feed-forward network without any kind of ensembling.
+            'plain',
+            #
+            # TabM
+            'tabm',
+            #
+            # TabM-mini
+            'tabm-mini',
+            #
+            # TabM-packed
+            'tabm-packed',
+            #
+            # TabM. The first adapter is initialized from the normal distribution.
+            # This variant was not used in the paper, but it may be useful in practice.
+            'tabm-normal',
+            #
+            # TabM-mini. The adapter is initialized from the normal distribution.
+            # This variant was not used in the paper.
+            'tabm-mini-normal',
+        ],
+        k: Union[None, int] = None,
+        share_training_batches: bool = True,
+    ) -> None:
+        # >>> Validate arguments.
+        assert n_num_features >= 0
+        assert n_num_features or cat_cardinalities
+        if arch_type == 'plain':
+            assert k is None
+            assert (
+                share_training_batches
+            ), 'If `arch_type` is set to "plain", then `simple` must remain True'
+        else:
+            assert k is not None
+            assert k > 0
+
+        super().__init__()
+
+        # >>> Continuous (numerical) features
+        first_adapter_sections = []  # See the comment in `_init_first_adapter`.
+
+        if n_num_features == 0:
+            assert bins is None
+            self.num_module = None
+            d_num = 0
+
+        elif num_embeddings is None:
+            assert bins is None
+            self.num_module = None
+            d_num = n_num_features
+            first_adapter_sections.extend(1 for _ in range(n_num_features))
+
+        else:
+            if bins is None:
+                self.num_module = make_module(
+                    **num_embeddings, n_features=n_num_features
+                )
+            else:
+                assert num_embeddings['type'].startswith('PiecewiseLinearEmbeddings')
+                self.num_module = make_module(**num_embeddings, bins=bins)
+            d_num = n_num_features * num_embeddings['d_embedding']
+            first_adapter_sections.extend(
+                num_embeddings['d_embedding'] for _ in range(n_num_features)
+            )
+
+        # >>> Categorical features
+        self.cat_module = (
+            OneHotEncoding0d(cat_cardinalities) if cat_cardinalities else None
+        )
+        first_adapter_sections.extend(cat_cardinalities)
+        d_cat = sum(cat_cardinalities)
+
+        # >>> Backbone
+        d_flat = d_num + d_cat
+        self.minimal_ensemble_adapter = None
+        # Any backbone can be here but we provide only MLP
+        self.backbone = make_module(d_in=d_flat, **backbone)
+
+        if arch_type != 'plain':
+            assert k is not None
+            first_adapter_init = (
+                None
+                if arch_type == 'tabm-packed'
+                else 'normal'
+                if arch_type in ('tabm-mini-normal', 'tabm-normal')
+                # For other arch_types, the initialization depends
+                # on the presence of num_embeddings.
+                else 'random-signs'
+                if num_embeddings is None
+                else 'normal'
+            )
+
+            if arch_type in ('tabm', 'tabm-normal'):
+                # Like BatchEnsemble, but all multiplicative adapters,
+                # except for the very first one, are initialized with ones.
+                assert first_adapter_init is not None
+                make_efficient_ensemble(
+                    self.backbone,
+                    LinearEfficientEnsemble,
+                    k=k,
+                    ensemble_scaling_in=True,
+                    ensemble_scaling_out=True,
+                    ensemble_bias=True,
+                    scaling_init='ones',
+                )
+                _init_first_adapter(
+                    _get_first_ensemble_layer(self.backbone).r,  # type: ignore[code]
+                    first_adapter_init,
+                    first_adapter_sections,
+                )
+
+            elif arch_type in ('tabm-mini', 'tabm-mini-normal'):
+                # MiniEnsemble
+                assert first_adapter_init is not None
+                self.minimal_ensemble_adapter = ScaleEnsemble(
+                    k,
+                    d_flat,
+                    init='random-signs' if num_embeddings is None else 'normal',
+                )
+                _init_first_adapter(
+                    self.minimal_ensemble_adapter.weight,  # type: ignore[code]
+                    first_adapter_init,
+                    first_adapter_sections,
+                )
+
+            elif arch_type == 'tabm-packed':
+                # Packed ensemble.
+                # In terms of the Packed Ensembles paper by Laurent et al.,
+                # TabM-packed is PackedEnsemble(alpha=k, M=k, gamma=1).
+                assert first_adapter_init is None
+                make_efficient_ensemble(self.backbone, NLinear, n=k)
+
+            else:
+                raise ValueError(f'Unknown arch_type: {arch_type}')
+
+        # >>> Output
+        d_block = backbone['d_block']
+        d_out = 1 if n_classes is None else n_classes
+        self.output = (
+            nn.Linear(d_block, d_out)
+            if arch_type == 'plain'
+            else NLinear(k, d_block, d_out)  # type: ignore[code]
+        )
+
+        # >>>
+        self.arch_type = arch_type
+        self.k = k
+        self.share_training_batches = share_training_batches
+
+    def forward(
+        self, x_num: Union[None, Tensor] = None, x_cat: Union[None, Tensor] = None
+    ) -> Tensor:
+        x = []
+        if x_num is not None:
+            x.append(x_num if self.num_module is None else self.num_module(x_num))
+        if x_cat is None:
+            assert self.cat_module is None
+        else:
+            assert self.cat_module is not None
+            x.append(self.cat_module(x_cat).float())
+        x = torch.column_stack([x_.flatten(1, -1) for x_ in x])
+
+        if self.k is not None:
+            if self.share_training_batches or not self.training:
+                # (B, D) -> (B, K, D)
+                x = x[:, None].expand(-1, self.k, -1)
+            else:
+                # (B * K, D) -> (B, K, D)
+                x = x.reshape(len(x) // self.k, self.k, *x.shape[1:])
+            if self.minimal_ensemble_adapter is not None:
+                x = self.minimal_ensemble_adapter(x)
+        else:
+            assert self.minimal_ensemble_adapter is None
+
+        x = self.backbone(x)
+        x = self.output(x)
+        if self.k is None:
+            # Adjust the output shape for plain networks to make them compatible
+            # with the rest of the script (loss, metrics, predictions, ...).
+            # (B, D_OUT) -> (B, 1, D_OUT)
+            x = x[:, None]
+        return x

autogluon/tabular/models/tabpfnmix/tabpfnmix_model.py

@@ -26,6 +26,8 @@ class TabPFNMixModel(AbstractModel):
 
     TabPFNMix is based off of the TabPFN and TabForestPFN models.
 
+    We recommend using Mitra instead, as it is an improved version of TabPFNMix.
+
     It is a tabular transformer model pre-trained on purely synthetic data.
 
     It currently has several limitations:
@@ -34,10 +36,13 @@ class TabPFNMixModel(AbstractModel):
     3. Does not support GPU
 
     For more information, refer to the `./_internals/README.md` file.
+
+    .. versionadded:: 1.2.0
     """
     ag_key = "TABPFNMIX"
     ag_name = "TabPFNMix"
     ag_priority = 45
+    seed_name = "random_state"
 
     weights_file_name = "model.pt"
 
@@ -119,6 +124,7 @@ class TabPFNMixModel(AbstractModel):
             raise AssertionError(f"Max allowed classes for the model is {max_classes}, " f"but found {self.num_classes} classes.")
 
         params = self._get_model_params()
+        random_state = params.pop(self.seed_name, self.default_random_seed)
         sample_rows = ag_params.get("sample_rows", None)
         sample_rows_val = ag_params.get("sample_rows_val", None)
         max_rows = ag_params.get("max_rows", None)
@@ -129,11 +135,11 @@ class TabPFNMixModel(AbstractModel):
 
         # TODO: Make sample_rows generic
         if sample_rows is not None and isinstance(sample_rows, int) and len(X) > sample_rows:
-            X, y = self._subsample_data(X=X, y=y, num_rows=sample_rows)
+            X, y = self._subsample_data(X=X, y=y, num_rows=sample_rows, random_state=random_state)
 
         # TODO: Make sample_rows generic
         if X_val is not None and y_val is not None and sample_rows_val is not None and isinstance(sample_rows_val, int) and len(X_val) > sample_rows_val:
-            X_val, y_val = self._subsample_data(X=X_val, y=y_val, num_rows=sample_rows_val)
+            X_val, y_val = self._subsample_data(X=X_val, y=y_val, num_rows=sample_rows_val, random_state=random_state)
 
         from ._internal.core.enums import Task
         if self.problem_type in [REGRESSION, QUANTILE]:
@@ -174,7 +180,7 @@ class TabPFNMixModel(AbstractModel):
         elif weights_path is not None:
             logger.log(15, f'\tLoading pre-trained weights from file... (weights_path="{weights_path}")')
 
-        cfg = ConfigRun(hyperparams=params, task=task, device=device)
+        cfg = ConfigRun(hyperparams=params, task=task, device=device, seed=random_state)
 
         if cfg.hyperparams["max_epochs"] == 0 and cfg.hyperparams["n_ensembles"] != 1:
             logger.log(
@@ -238,7 +244,7 @@ class TabPFNMixModel(AbstractModel):
         return self
 
     # TODO: Make this generic by creating a generic `preprocess_train` and putting this logic prior to `_preprocess`.
-    def _subsample_data(self, X: pd.DataFrame, y: pd.Series, num_rows: int, random_state=0) -> (pd.DataFrame, pd.Series):
+    def _subsample_data(self, X: pd.DataFrame, y: pd.Series, num_rows: int, random_state: int | None = 0) -> (pd.DataFrame, pd.Series):
         num_rows_to_drop = len(X) - num_rows
         X, _, y, _ = generate_train_test_split(
             X=X,
@@ -311,11 +317,11 @@ class TabPFNMixModel(AbstractModel):
 
     def _get_maximum_resources(self) -> dict[str, int | float]:
         # torch model trains slower when utilizing virtual cores and this issue scale up when the number of cpu cores increases
-        return {"num_cpus": ResourceManager.get_cpu_count_psutil(logical=False)}
+        return {"num_cpus": ResourceManager.get_cpu_count(only_physical_cores=True)}
 
     def _get_default_resources(self) -> tuple[int, float]:
-        # logical=False is faster in training
-        num_cpus = ResourceManager.get_cpu_count_psutil(logical=False)
+        # only_physical_cores=True is faster in training
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
         num_gpus = 0
         return num_cpus, num_gpus