From c0c7c195b10d99ca3f0e4f1253e06262e71b384b Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Thu, 28 Sep 2023 11:46:30 +0200 Subject: [PATCH 1/5] clean up global flags HAS*_TORCH * in nn module use HAS_TORCH from equisolve/__init__.py * mv HAS_METATENSOR_TORCH to equisolve/__init__.py * add function refresh_global_flags that allows to refresh the global flags --- src/equisolve/__init__.py | 37 ++++++++++++++++++++++++------- src/equisolve/nn/__init__.py | 6 +---- src/equisolve/nn/module_tensor.py | 11 +++++---- 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/src/equisolve/__init__.py b/src/equisolve/__init__.py index 1fde0a6..1d973a9 100644 --- a/src/equisolve/__init__.py +++ b/src/equisolve/__init__.py @@ -9,11 +9,32 @@ __version__ = "0.0.0-dev" __authors__ = "the equisolve development team" -# For a global consistent state of the package, we try to load here torch, -# since torch is an optional dependency -try: - import torch # noqa: F401 - - HAS_TORCH = True -except ImportError: - HAS_TORCH = False + +def refresh_global_flags(): + """ + Refreshes all global flags set on import of library. This function might be useful + if one is in an interactive session and installed some of the optional dependenicies + (torch, metatensor-torch) after importing the library. + """ + global HAS_TORCH + global HAS_METATENSOR_TORCH + + try: + import torch # noqa: F401 + + HAS_TORCH = True + except ImportError: + HAS_TORCH = False + + try: + from metatensor.torch import Labels, TensorBlock, TensorMap # noqa: F401 + + HAS_METATENSOR_TORCH = True + except ImportError: + from metatensor import Labels, TensorBlock, TensorMap # noqa: F401 + + HAS_METATENSOR_TORCH = False + + +# For a global consistent state of the package, we set the global flags once here +refresh_global_flags() diff --git a/src/equisolve/nn/__init__.py b/src/equisolve/nn/__init__.py index b29d4c8..38c51a3 100644 --- a/src/equisolve/nn/__init__.py +++ b/src/equisolve/nn/__init__.py @@ -1,9 +1,5 @@ -try: - import torch # noqa: F401 +from .. import HAS_TORCH - HAS_TORCH = True -except ImportError: - HAS_TORCH = False if HAS_TORCH: from .module_tensor import Linear, ModuleTensorMap # noqa: F401 diff --git a/src/equisolve/nn/module_tensor.py b/src/equisolve/nn/module_tensor.py index e06ea3e..2826f41 100644 --- a/src/equisolve/nn/module_tensor.py +++ b/src/equisolve/nn/module_tensor.py @@ -1,11 +1,10 @@ -try: - from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap +from .. import HAS_METATENSOR_TORCH - HAS_METATENSOR_TORCH = True -except ImportError: - from metatensor import Labels, LabelsEntry, TensorBlock, TensorMap - HAS_METATENSOR_TORCH = False +if HAS_METATENSOR_TORCH: + from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap +else: + from metatensor import Labels, LabelsEntry, TensorBlock, TensorMap from copy import deepcopy from typing import List, Optional From 177b4efa44380a7bed2b0e923c8413364e4627a9 Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Thu, 28 Sep 2023 12:01:03 +0200 Subject: [PATCH 2/5] rm temporary tensor map to dict solution tensor maps are now pickable so we dont need to do it anymore --- src/equisolve/numpy/models/linear_model.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/equisolve/numpy/models/linear_model.py b/src/equisolve/numpy/models/linear_model.py index 630355b..caedd70 100644 --- a/src/equisolve/numpy/models/linear_model.py +++ b/src/equisolve/numpy/models/linear_model.py @@ -16,7 +16,7 @@ from ... import HAS_TORCH from ...module import NumpyModule, _Estimator from ...utils.metrics import rmse -from ..utils import array_from_block, dict_to_tensor_map, tensor_map_to_dict +from ..utils import array_from_block class _Ridge(_Estimator): @@ -307,8 +307,7 @@ def fit( weights_blocks.append(weight_block) - # convert weights to a dictionary allowing pickle dump of an instance - self._weights = tensor_map_to_dict(TensorMap(X.keys, weights_blocks)) + self._weights = TensorMap(X.keys, weights_blocks) return self @@ -319,7 +318,7 @@ def weights(self) -> TensorMap: if self._weights is None: raise ValueError("No weights. Call fit method first.") - return dict_to_tensor_map(self._weights) + return self._weights def predict(self, X: TensorMap) -> TensorMap: """ From 29799cbb658cd5065e23e2b6e2650f709034fc6c Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Thu, 28 Sep 2023 12:45:19 +0200 Subject: [PATCH 3/5] add support for TensorMap as argument for bias --- src/equisolve/nn/module_tensor.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/equisolve/nn/module_tensor.py b/src/equisolve/nn/module_tensor.py index 2826f41..b7b7950 100644 --- a/src/equisolve/nn/module_tensor.py +++ b/src/equisolve/nn/module_tensor.py @@ -7,7 +7,7 @@ from metatensor import Labels, LabelsEntry, TensorBlock, TensorMap from copy import deepcopy -from typing import List, Optional +from typing import List, Optional, Union import torch from torch.nn import Module, ModuleDict @@ -167,15 +167,28 @@ class Linear(ModuleTensorMap): properties, the labels of the properties cannot be persevered. :param bias: - See :py:class:`torch.nn.Linear` + See :py:class:`torch.nn.Linear` for bool as input. For each TensorMap key the + bias can be also individually tuend by using a TensorMap with one value for the + bool. """ def __init__( self, in_tensor: TensorMap, out_tensor: TensorMap, - bias: bool = True, + bias: Union[bool, TensorMap] = True, ): + if isinstance(bias, bool): + blocks = [ + TensorBlock( + values=torch.tensor(bias).reshape(1, 1), + samples=Labels.range("_", 1), + components=[], + properties=Labels.range("_", 1), + ) + for _ in in_tensor.keys + ] + bias = TensorMap(keys=in_tensor.keys, blocks=blocks) module_map = ModuleDict() for key, in_block in in_tensor.items(): module_key = ModuleTensorMap.module_key(key) @@ -183,7 +196,7 @@ def __init__( module = torch.nn.Linear( len(in_block.properties), len(out_block.properties), - bias, + bias.block(key).values.flatten()[0], in_block.values.device, in_block.values.dtype, ) From 75a997c6245b375d5528c9c801b6d174b73e242b Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Mon, 2 Oct 2023 13:51:59 +0200 Subject: [PATCH 4/5] remove NumpyRidge TorchRidge this commit prepares the Ridge class to use an export function to obtain TorchScriptable modules --- src/equisolve/numpy/models/linear_model.py | 24 ++-------------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/src/equisolve/numpy/models/linear_model.py b/src/equisolve/numpy/models/linear_model.py index caedd70..a36e3ad 100644 --- a/src/equisolve/numpy/models/linear_model.py +++ b/src/equisolve/numpy/models/linear_model.py @@ -13,13 +13,12 @@ import scipy.linalg from metatensor import Labels, TensorBlock, TensorMap -from ... import HAS_TORCH -from ...module import NumpyModule, _Estimator +from ...module import _Estimator from ...utils.metrics import rmse from ..utils import array_from_block -class _Ridge(_Estimator): +class Ridge(_Estimator): r"""Linear least squares with l2 regularization for :class:`metatensor.Tensormap`'s. Weights :math:`w` are calculated according to @@ -351,22 +350,3 @@ def score(self, X: TensorMap, y: TensorMap, parameter_key: str) -> float: """ y_pred = self.predict(X) return rmse(y, y_pred, parameter_key) - - -class NumpyRidge(_Ridge, NumpyModule): - def __init__(self) -> None: - NumpyModule.__init__(self) - _Ridge.__init__(self) - - -if HAS_TORCH: - import torch - - class TorchRidge(_Ridge, torch.nn.Module): - def __init__(self) -> None: - torch.nn.Module.__init__(self) - _Ridge.__init__(self) - - Ridge = TorchRidge -else: - Ridge = NumpyRidge From 1331a29746b857944e6c1d14fff2d7f483184073 Mon Sep 17 00:00:00 2001 From: Alexander Goscinski Date: Thu, 28 Sep 2023 13:52:18 +0200 Subject: [PATCH 5/5] add export_torch_module to linear module and tests * add from_weight constructor to equisolve Linear torch module * add helper function core_tensor_map_to_torch and transpose_tensor_map --- src/equisolve/nn/module_tensor.py | 28 ++++++ src/equisolve/numpy/models/linear_model.py | 31 ++++++- src/equisolve/numpy/utils.py | 88 ++++++++++++++++++- .../numpy/models/linear_model.py | 28 ++++++ 4 files changed, 173 insertions(+), 2 deletions(-) diff --git a/src/equisolve/nn/module_tensor.py b/src/equisolve/nn/module_tensor.py index b7b7950..fd38ecf 100644 --- a/src/equisolve/nn/module_tensor.py +++ b/src/equisolve/nn/module_tensor.py @@ -242,6 +242,34 @@ def from_module( module = torch.nn.Linear(in_features, out_features, bias, device, dtype) return ModuleTensorMap.from_module(in_keys, module, many_to_one, out_tensor) + @classmethod + def from_weights(cls, weights: TensorMap, bias: Optional[TensorMap] = None): + """ + :param weights: + The weight tensor map from which we create the linear modules. The + properties of the tensor map describe the input dimension and the samples + describe the output dimension. + + :param bias: + The weight tensor map from which we create the linear layers. + """ + module_map = ModuleDict() + for key, weights_block in weights.items(): + module_key = ModuleTensorMap.module_key(key) + module = torch.nn.Linear( + len(weights_block.samples), + len(weights_block.properties), + bias=False, + device=weights_block.values.device, + dtype=weights_block.values.dtype, + ) + module.weight = torch.nn.Parameter(weights_block.values.T) + if bias is not None: + module.bias = torch.nn.Parameter(bias.block(key).values) + module_map[module_key] = module + + return ModuleTensorMap(module_map, weights) + def forward(self, tensor: TensorMap) -> TensorMap: # added to appear in doc, :inherited-members: is not compatible with torch return super().forward(tensor) diff --git a/src/equisolve/numpy/models/linear_model.py b/src/equisolve/numpy/models/linear_model.py index a36e3ad..8c997a0 100644 --- a/src/equisolve/numpy/models/linear_model.py +++ b/src/equisolve/numpy/models/linear_model.py @@ -15,7 +15,7 @@ from ...module import _Estimator from ...utils.metrics import rmse -from ..utils import array_from_block +from ..utils import array_from_block, core_tensor_map_to_torch, transpose_tensor_map class Ridge(_Estimator): @@ -350,3 +350,32 @@ def score(self, X: TensorMap, y: TensorMap, parameter_key: str) -> float: """ y_pred = self.predict(X) return rmse(y, y_pred, parameter_key) + + def export_torch_module(self, device=None, dtype=None): + """ + Export existing weights to a child class :py:class:`torch.nn.Module` so it can + :py:mod:`torch.jit` utils can be applied. + + :param device: + :py:class:`torch.device` of values in the resulting module + + :param dtye: + :py:class:`torch.dtype` of the values in the resulting module + + :returns linear: + a :py:class:`equisolve.nn.Linear` + """ + from ... import HAS_METATENSOR_TORCH + + if not HAS_METATENSOR_TORCH: + raise ImportError( + "To export your model to TorchScript torch needs to be installed. " + "Please install torch, then reimport equisolve or " + "use equisolve.refresh_global_flags()." + ) + from ...nn import Linear + + torch_weights = core_tensor_map_to_torch( + transpose_tensor_map(self.weights), device, dtype + ) + return Linear.from_weights(torch_weights) diff --git a/src/equisolve/numpy/utils.py b/src/equisolve/numpy/utils.py index e219a67..cfaddc7 100644 --- a/src/equisolve/numpy/utils.py +++ b/src/equisolve/numpy/utils.py @@ -11,7 +11,7 @@ import metatensor import numpy as np -from metatensor import TensorBlock, TensorMap +from metatensor import Labels, TensorBlock, TensorMap def array_from_block(block: TensorBlock) -> np.ndarray: @@ -76,3 +76,89 @@ def dict_to_tensor_map(tensor_map_dict: dict): tmp_filename = tempfile.mktemp() + ".npz" np.savez(tmp_filename, **tensor_map_dict) return metatensor.load(tmp_filename) + + +def core_tensor_map_to_torch(core_tensor: TensorMap, device=None, dtype=None): + """Transforms a tensor map from metatensor-core to metatensor-torch + + :param core_tensor: + tensor map from metatensor-core + + :param device: + :py:class:`torch.device` of values in the resulting tensor map + + :param dtye: + :py:class:`torch.dtype` of values in the resulting tensor map + + :returns torch_tensor: + tensor map from metatensor-torch + """ + from metatensor.torch import TensorMap as TorchTensorMap + + torch_blocks = [] + for _, core_block in core_tensor.items(): + torch_blocks.append(core_tensor_block_to_torch(core_block, device, dtype)) + torch_keys = core_labels_to_torch(core_tensor.keys) + return TorchTensorMap(torch_keys, torch_blocks) + + +def core_tensor_block_to_torch(core_block: TensorBlock, device=None, dtype=None): + """Transforms a tensor block from metatensor-core to metatensor-torch + + :param core_block: + tensor block from metatensor-core + + :param device: + :py:class:`torch.device` of values in the resulting block and labels + + :param dtye: + :py:class:`torch.dtype` of values in the resulting block and labels + + :returns torch_block: + tensor block from metatensor-torch + """ + import torch + from metatensor.torch import TensorBlock as TorchTensorBlock + + return TorchTensorBlock( + values=torch.tensor(core_block.values, device=device, dtype=dtype), + samples=core_labels_to_torch(core_block.samples, device=device), + components=[ + core_labels_to_torch(component, device=device) + for component in core_block.components + ], + properties=core_labels_to_torch(core_block.properties, device=device), + ) + + +def core_labels_to_torch(core_labels: Labels, device=None): + """Transforms labels from metatensor-core to metatensor-torch + + :param core_block: + tensor block from metatensor-core + + :param device: + :py:class:`torch.device` of values in the resulting labels + + :returns torch_block: + labels from metatensor-torch + """ + import torch + from metatensor.torch import Labels as TorchLabels + + return TorchLabels( + core_labels.names, torch.tensor(core_labels.values, device=device) + ) + + +def transpose_tensor_map(tensor: TensorMap): + blocks = [] + for block in tensor.blocks(): + block = TensorBlock( + values=block.values.T, + samples=block.properties, + components=block.components, + properties=block.samples, + ) + blocks.append(block) + return TensorMap(tensor.keys, blocks) diff --git a/tests/equisolve_tests/numpy/models/linear_model.py b/tests/equisolve_tests/numpy/models/linear_model.py index e172b6c..16c96a3 100644 --- a/tests/equisolve_tests/numpy/models/linear_model.py +++ b/tests/equisolve_tests/numpy/models/linear_model.py @@ -11,7 +11,9 @@ from metatensor import Labels, TensorBlock, TensorMap from numpy.testing import assert_allclose, assert_equal +from equisolve import HAS_METATENSOR_TORCH from equisolve.numpy.models import Ridge +from equisolve.numpy.utils import core_tensor_map_to_torch from ..utilities import tensor_to_tensormap @@ -79,6 +81,32 @@ def equisolve_solver_from_numpy_arrays( clf.fit(X=X, y=y, alpha=alpha, sample_weight=sw, solver=solver) return clf + @pytest.mark.skipif( + not (HAS_METATENSOR_TORCH), reason="requires metatensor-torch to be run" + ) + def test_export_torch_module(self): + """Test if ridge is working and all shapes are converted correctly. + Test is performed for two blocks. + """ + + num_targets = 50 + num_properties = 5 + + # Create input values + X_arr = self.rng.random([2, num_targets, num_properties]) + y_arr = self.rng.random([2, num_targets, 1]) + + X = tensor_to_tensormap(X_arr) + y = tensor_to_tensormap(y_arr) + + clf = Ridge() + clf.fit(X=X, y=y) + y_pred_torch = core_tensor_map_to_torch(clf.predict(X)) + + module = clf.export_torch_module() + y_pred_torch_module = module.forward(core_tensor_map_to_torch(X)) + metatensor.torch.allclose_raise(y_pred_torch, y_pred_torch_module) + num_properties = np.array([91]) num_targets = np.array([1000]) means = np.array([-0.5, 0, 0.1])