Comments (3)
Nice spot. Sorry, this is a bug in the TF code introduced when porting parallelisation over multiple GPUs from TF-Replicator to Distribution Strategy. The results in the FermiNet paper used TF-Replicator, which has a slightly different API to Distribution Strategy. With TF-Replicator, we are using the concatenated distribution on each GPU. The line you highlighted should read
concat_data = tf.concat([self.hf_data_gen.walkers, self.data_gen.walkers], axis=1)
We made a couple of different design decisions in the JAX version. We are currently only using the FermiNet as the distribution in pretraining. I think the pretraining distribution is not super critical -- the most important thing is to train the network such that it's closer to the ground state wavefunction and that the determinants aren't extremely low rank. Without pretraining, the network can start with an energy in the 1000s of hartrees (positive!) and optimising the network through that energy landscape is ... let's say painful. Certainly we don't see any noticeable difference in training FermiNet after pretraining in the TF or JAX codes.
from ferminet.
Got it. Thanks a lot for the detailed explanation!
from ferminet.
Whilst we don't believe this makes a significant difference, the TF version on master is now updated to sample from both distributions during pretraining.
from ferminet.
Related Issues (20)
- Jax install - issue with correct version number HOT 1
- AttributeError: module 'jax.core' has no attribute 'extract_call_jaxpr' HOT 1
- Jax error running on A100 GPU (everything is okay on CPU) HOT 2
- unable to setup HOT 1
- The proper way to cite FermiNet repo HOT 1
- Ground State Energies HOT 2
- Question about pbc ewald part. HOT 2
- nan when training with 'adam' HOT 1
- About configs HOT 3
- Question About load Checkpoint HOT 1
- Evaluating logprob using batch_network in train HOT 1
- Issue on running pytest HOT 5
- Extension of PBC code to 1D HOT 7
- Something went wrong in RepeatedDenseBlock.update_curvature_matrix_estimate HOT 2
- Different results obtained from the paper for ch3nh2 HOT 2
- kfac_jax error when running H2 example script HOT 2
- Upstream breaking change in `kfac-jax`
- KeyError raised after burn-in MCMC steps HOT 1
- Logdet Bug Similar to e9f8c64 HOT 1
- Natural Excited States to nan HOT 1
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from ferminet.