add neural network modules that work with tensor maps #62

agoscinski · 2023-09-14T18:48:02Z

From the discussion in #13. Should replace https://github.com/bananenpampe/H2O/blob/move-rascaline/model/nn/modules.py

The goal is to offer a flexible wrapper class around torch modules that are applied on each block of a input TensorMap while also offer functions to wrap simpler cases, for example when the same module is used for all blocks (for that there is the from_module constructor).

What is missing, is a functionality that allows to merge multiple blocks into a new one (e.g. in species compression), but I also don't think it is a good idea to put this in the same class.

TODO (will add these after first review):

add another test for using the __init__ of Linear
add another test for using the from_module of ModuleTensorMap
move equisolve.nn to equisolve.torch.nn (needs to be discussed in meeting)
rascaline has been updated Update metatensor metatensor/featomic#232 we can remove the requirements.txt (maybe keeping it is smarter because of cross linking of packages with sphinx, need to check this.)

📚 Documentation preview 📚: https://equisolve--62.org.readthedocs.build/en/62/

PicoCentauri · 2023-09-19T11:58:10Z

What is missing, is a functionality that allows to merge multiple blocks into a new one (e.g. in species compression), but I also don't think it is a good idea to put this in the same class.

This can be done by the user using vanilla tensormap functions, right?

Luthaf

The code looks good, the documentation will still need some work: how to wrap custom module inside equisolve.nn.Module, how/when to use these/... but this can come at a later date.

Luthaf · 2023-09-19T12:10:45Z

src/equisolve/nn/__init__.py

+    HAS_TORCH = False
+
+if HAS_TORCH:
+    from .module_tensor import LinearTensorMap, ModuleTensorMap  # noqa: F401


naming wise, I would call these Linear and Module, and users can then use torch.nn.Linear/torch.nn.Module or equisolve.nn.Linear/equisolve.nn.Module for the different behavior.

okay lets try this. We might move these equisolve.torch.nn.Linear, to better distinguish the methods here from the once supporting also numpy backend. But I would do the refactor in another PR.

PicoCentauri · 2023-09-19T12:03:34Z

.github/workflows/tests.yml

+    - name: run Python metatensor-core tests
+      run: tox -e metatensor-core-numpy-tests
+
+    - name: run Python metatensor-core tests
+      run: tox -e metatensor-core-torch-tests
+
+    - name: run Python metatensor-torch tests
+      run: tox -e metatensor-torch-tests


I think you can remove the metatensor from the name. We don't have tests without metatensor I assume.

PicoCentauri · 2023-09-19T12:03:56Z

docs/requirements.txt

@@ -10,3 +10,7 @@ furo
 rascaline @ https://github.com/luthaf/rascaline/archive/fb5332f.zip
 sphinx >=4.4
 sphinx-gallery
+# temporary add torch to dependencies until rascaline is on metatensor
+# version 3c5fee40


Suggested change

# version 3c5fee40

# version 3c5fee40

I removed the comment, because we need to have a torch dependency anyway later for the examples. The dependency is now given through extra requirement on metatensor-torch in tox.ini and .readthedocs.yml

PicoCentauri · 2023-09-19T12:22:13Z

src/equisolve/nn/module_tensor.py

+
+        return TensorMap(tensor.keys, out_blocks)
+
+    def forward_block(self, key: LabelsEntry, block: TensorBlock) -> TensorBlock:


I think this can also be a Labels object right?

if you get it by iterating over the keys, you'll get LabelsEntry

The input is one key of a TensorMap labels object. It has a bit different structure than metatensor-operations, because here we need the key to actually figure which module to apply on the block

PicoCentauri · 2023-09-19T12:25:28Z

src/equisolve/nn/module_tensor.py

+    :param out_tensor: a tensor map that has the output properties of the Labels
+    """
+
+    def __init__(self, module_map: ModuleDict, out_tensor: Optional[TensorMap] = None):


I am not sure what the out_tensor is doing. Is this the TensorMap where the results will be written in?

added better documentation. Also there was a bug how it was used in the Linear module

to wrap existing torch arbitrary modules and apply them on each key of a tensor map

agoscinski

What is missing, is a functionality that allows to merge multiple blocks into a new one (e.g. in species compression), but I also don't think it is a good idea to put this in the same class.

This can be done by the user using vanilla tensormap functions, right?

Yes, I am not sure if one can create another Module for this purpose that is generic enough to be useful for multiple applications without being too abstract and less efficient than doing it "manually". For example the species decomposition in alchemical learning. You are transforming the species neighbor sparsity to a new property the pseudo species. So this combines a movement to properties and a linear combination. Don't have good ideas here.

The code looks good, the documentation will still need some work: how to wrap custom module inside equisolve.nn.Module, how/when to use these/... but this can come at a later date.

Improved the documentation. I will do an issue for an example for the neural network.

EDIT: figuring out how to add PIP_EXTRA_INDEX_URL to readthedocs, otherwise everything passes
EDIT: seems to work now, but I reached concurrency limit^^

agoscinski · 2023-09-21T20:23:30Z

src/equisolve/nn/module_tensor.py

+    :param out_tensor: a tensor map that has the output properties of the Labels
+    """
+
+    def __init__(self, module_map: ModuleDict, out_tensor: Optional[TensorMap] = None):


added better documentation. Also there was a bug how it was used in the Linear module

agoscinski · 2023-09-21T20:26:23Z

src/equisolve/nn/module_tensor.py

+
+        return TensorMap(tensor.keys, out_blocks)
+
+    def forward_block(self, key: LabelsEntry, block: TensorBlock) -> TensorBlock:


The input is one key of a TensorMap labels object. It has a bit different structure than metatensor-operations, because here we need the key to actually figure which module to apply on the block

agoscinski · 2023-09-21T20:31:08Z

docs/requirements.txt

@@ -10,3 +10,7 @@ furo
 rascaline @ https://github.com/luthaf/rascaline/archive/fb5332f.zip
 sphinx >=4.4
 sphinx-gallery
+# temporary add torch to dependencies until rascaline is on metatensor
+# version 3c5fee40


I removed the comment, because we need to have a torch dependency anyway later for the examples. The dependency is now given through extra requirement on metatensor-torch in tox.ini and .readthedocs.yml

agoscinski · 2023-09-21T20:31:45Z

.github/workflows/tests.yml

+    - name: run Python metatensor-core tests
+      run: tox -e metatensor-core-numpy-tests
+
+    - name: run Python metatensor-core tests
+      run: tox -e metatensor-core-torch-tests
+
+    - name: run Python metatensor-torch tests
+      run: tox -e metatensor-torch-tests


agoscinski · 2023-09-21T20:32:03Z

src/equisolve/nn/__init__.py

+    HAS_TORCH = False
+
+if HAS_TORCH:
+    from .module_tensor import LinearTensorMap, ModuleTensorMap  # noqa: F401


okay lets try this. We might move these equisolve.torch.nn.Linear, to better distinguish the methods here from the once supporting also numpy backend. But I would do the refactor in another PR.

the existing tests are moved to core-numpy-tests and new environments for metatensor-core with torch metatensor-torch are added

add tests for from_module and default constructor

the function random_single_block_no_components_tensor_map is now in one global utilities

add tests for from_module and default constructor

this optional dependency loads metatensor-torch

This properly references TensorMap types in the doc instead of only showing abstruse ScriptClass types

PicoCentauri

Thanks for the improced Docs @agoscinski !

agoscinski requested a review from bananenpampe September 14, 2023 18:48

agoscinski force-pushed the tensormap-module branch from a27c00c to d5ec1ae Compare September 15, 2023 08:02

agoscinski mentioned this pull request Sep 15, 2023

torch.Tensor support for methods in numpy submodule #63

Open

agoscinski force-pushed the tensormap-module branch 3 times, most recently from f01e400 to 665688d Compare September 15, 2023 11:46

agoscinski marked this pull request as ready for review September 15, 2023 12:03

agoscinski force-pushed the tensormap-module branch from 8e8ad63 to b3e8754 Compare September 15, 2023 14:11

agoscinski requested a review from Luthaf September 15, 2023 14:11

agoscinski changed the title ~~add initial modules that work with tensor maps~~ add neural network modules that work with tensor maps Sep 15, 2023

agoscinski requested a review from PicoCentauri September 19, 2023 10:55

Luthaf reviewed Sep 19, 2023

View reviewed changes

PicoCentauri reviewed Sep 19, 2023

View reviewed changes

add initial module for tensor maps

1aef4c4

to wrap existing torch arbitrary modules and apply them on each key of a tensor map

agoscinski force-pushed the tensormap-module branch from b3e8754 to 32951b5 Compare September 22, 2023 12:45

agoscinski commented Sep 22, 2023

View reviewed changes

agoscinski added 9 commits September 22, 2023 14:54

add tests for metatensor-torch subpackage

0fdeeab

the existing tests are moved to core-numpy-tests and new environments for metatensor-core with torch metatensor-torch are added

add tests for equisolve.nn.Linear

1f35ff0

add tests for from_module and default constructor

merge tensor creation in numpy and nn utilities

f2c9533

the function random_single_block_no_components_tensor_map is now in one global utilities

add tests for ModuleTensorMap

a804440

add tests for from_module and default constructor

update to a recent rascaline+metatensor version

b9fdd94

add torch optional dependencies to equisolve

4f08ce0

this optional dependency loads metatensor-torch

add extra torch to doc requirements

83f01c2

replace metatensor-torch in deps with extras torch

79348e6

add torch to intersphinx_mapping

4f7af24

agoscinski force-pushed the tensormap-module branch from 32951b5 to 5f732e7 Compare September 22, 2023 12:54

agoscinski requested review from PicoCentauri and Luthaf September 22, 2023 13:19

agoscinski force-pushed the tensormap-module branch from 5f732e7 to 5aa8936 Compare September 22, 2023 14:27

agoscinski mentioned this pull request Sep 22, 2023

Example using ModuleTensorMap #65

Open

agoscinski added 2 commits September 22, 2023 16:34

set environment var METATENSOR_IMPORT_FOR_SPHINX

bd37e0c

This properly references TensorMap types in the doc instead of only showing abstruse ScriptClass types

add neural network module to the doc

5388061

agoscinski force-pushed the tensormap-module branch from 5aa8936 to 5388061 Compare September 22, 2023 14:34

PicoCentauri approved these changes Sep 26, 2023

View reviewed changes

agoscinski merged commit c4ae847 into main Sep 26, 2023
8 checks passed

agoscinski deleted the tensormap-module branch September 26, 2023 15:14

agoscinski mentioned this pull request Dec 2, 2023

Implementation of ModuleMap in metatensor-learn metatensor/metatensor#427

Merged

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add neural network modules that work with tensor maps #62

add neural network modules that work with tensor maps #62

agoscinski commented Sep 14, 2023 •

edited

Loading

PicoCentauri commented Sep 19, 2023

Luthaf left a comment

Luthaf Sep 19, 2023

agoscinski Sep 21, 2023

PicoCentauri Sep 19, 2023

agoscinski Sep 21, 2023

PicoCentauri Sep 19, 2023

agoscinski Sep 21, 2023

PicoCentauri Sep 19, 2023

Luthaf Sep 19, 2023

agoscinski Sep 21, 2023 •

edited

Loading

PicoCentauri Sep 19, 2023

agoscinski Sep 21, 2023

agoscinski left a comment •

edited

Loading

agoscinski Sep 21, 2023

agoscinski Sep 21, 2023 •

edited

Loading

agoscinski Sep 21, 2023

agoscinski Sep 21, 2023

agoscinski Sep 21, 2023

PicoCentauri left a comment


		return TensorMap(tensor.keys, out_blocks)

		def forward_block(self, key: LabelsEntry, block: TensorBlock) -> TensorBlock:

add neural network modules that work with tensor maps #62

add neural network modules that work with tensor maps #62

Conversation

agoscinski commented Sep 14, 2023 • edited Loading

PicoCentauri commented Sep 19, 2023

Luthaf left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

agoscinski Sep 21, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

agoscinski left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

agoscinski Sep 21, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

PicoCentauri left a comment

Choose a reason for hiding this comment

agoscinski commented Sep 14, 2023 •

edited

Loading

agoscinski Sep 21, 2023 •

edited

Loading

agoscinski left a comment •

edited

Loading

agoscinski Sep 21, 2023 •

edited

Loading