Skip to content
This repository has been archived by the owner on Dec 5, 2024. It is now read-only.

Example notebook with PMAP sampling of RBMs trained on MNIST digits #80

Merged
merged 8 commits into from
Oct 17, 2021

Conversation

StannisZhou
Copy link
Contributor

@StannisZhou StannisZhou commented Oct 14, 2021

Resolves the first todo in #77. Also includes a simple change with which we can index 1D variable arrays with integers (instead of length 1 tuples).

Currently a lot of time is spent on graph construction (adding factors and constructing factor wirings) since PGMax is not making use of any of the structures present in an RBM.

In the future, to improve usability, it would be helpful to:

  1. See if we can find an efficient way to parallelize many of the involved for loops (Optimize and parallelize structure compiling function #18)
  2. Introduce a mechanism to serialize a model to disk
  3. Improve interface so that we can explicitly use already constructed wiring.

Separately, it would be helpful to introduce a pure functional interface so that we can leverage functionalities like vmap in JAX (which would allow us to parallelize sampling multiple digits in this use case). Connected to #68

See below for some example samples
Screenshot from 2021-10-13 21-22-45
Screenshot from 2021-10-12 23-33-37
:

Copy link
Contributor

@NishanthJKumar NishanthJKumar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks mostly good, but 1 potentially-important comment about dependencies and another more minor comment about testing. Will approve after these are addressed/replied to!

examples/rbm.py Outdated Show resolved Hide resolved
tests/fg/test_groups.py Outdated Show resolved Hide resolved
Copy link
Contributor

@NishanthJKumar NishanthJKumar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks for addressing the previous comments so quickly 😄

@StannisZhou
Copy link
Contributor Author

Looks like the use of __file__ would break the notebooks. Having the notebooks running seems more important so I got rid of the __file__ and also excluded heretic model example from unit tests for now.

@StannisZhou StannisZhou merged commit 5bad82c into vicariousinc:master Oct 17, 2021
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants