Skip to content

Allow pure Python-based Machines and PyTorch modules

Vicentini Filippo requested to merge py_machine into v2.1

Created by: gcarleo

Goal of this PR is mainly to do the changes necessary to allow the use of generic (pure python) machines (without the need to derive them from the C++ netket machines using trampoline python classes and other nasty stuff).

This PR provides the following main changes:

  1. new local_values in python, with basically the same performance as the old one written in C++.

  2. get_conn has been rewritten to work in batches, and also addressing #316 (closed).

  3. An AbstractMachine prototype in python. @noamwies please check if this fits with your needs.

  4. There is a pure numpy (with a little C++ help for log_cosh...) implementing a spin RBM (PyRbm, largely inspired by @twesterhout's work). Notice that this does not derive from the C++ class as the previous implementation, and it can work as it is in all the samplers we have (since the samplers switch automatically to the python-based implementation). As of now, PyRbm performance on small systems (say 20 spins, alpha=1) is a factor of 3 worse than the C++. This can be improved further (we should get to about x 1.5 C++ version on small machines and on par for larger ones) removing some bottlenecks that we have in the python samplers (I will do that in a later PR).

  5. Partial, but working, support for generic pytorch models. Performance is not great on small systems, but in my tests with the current setup already around 100 spins pytorch on CPU outperforms our C++ implementation. Leveraging better sampling techniques and/or using some of the threading tricks only @twesterhout knows, we could further shrink the gap on small systems, but overall I think this is pretty good at this stage. We are sorely missing a clean and efficient implementation of the Jacobian in PyTorch, however there is some work going on there, so hopefully we will have it at some point this year!

Merge request reports