adding complete forward and backward pass for ferminet pretraining #3553

shaipranesh2 · 2023-09-03T19:41:26Z

Description

Fix #(issue)

Type of change

Please check the option that is related to your PR.

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
- In this case, we recommend to discuss your modification on GitHub issues before creating the PR
Documentations (modification for documents)

Checklist

…into ferminet_hf

shaipranesh2 · 2023-09-27T15:20:22Z

deepchem/models/torch_models/ferminet.py

@@ -71,6 +65,8 @@ def __init__(self,

        Attributes
        ----------
+        running_diff: torch.Tensor


Adding an attribute tensor to keep track of running sum of MSELoss

shaipranesh2 · 2023-09-27T15:20:50Z

deepchem/models/torch_models/ferminet.py

@@ -148,6 +145,30 @@ def forward(self, input) -> torch.Tensor:
            0].forward(one_electron, one_electron_vector_permuted)
        return psi

+    def loss(self,


loss functions, right now its implemented for pertaining (MSELoss)

shaipranesh2 · 2023-09-27T15:21:13Z

deepchem/models/torch_models/ferminet.py

@@ -205,10 +226,8 @@ def __init__(self,
        Torch tensor containing electrons for each atom in the nucleus
    molecule: ElectronSampler
        ElectronSampler object which performs MCMC and samples electrons
-    loss_value: Optional[torch.Tensor] (default None)
+    loss_value: torch.Tensor (default torch.tensor(0))


attribute that keeps track of the last loss value. It differs form running mean by calculating the mean over all batches and the number of steps in MCMC

shaipranesh2 · 2023-09-27T15:22:30Z

deepchem/models/torch_models/ferminet.py


        self.molecule: ElectronSampler = ElectronSampler(
            batch_no=self.batch_no,
            central_value=self.nucleon_pos,
            seed=self.seed,
-            f=lambda x: test_f(x),  # Will be replaced in successive PR
+            f=lambda x: self.random_walk(x


changing the sampling function to the random walk function which returns the log wavefucntion probability

shaipranesh2 · 2023-09-27T15:23:24Z

deepchem/models/torch_models/ferminet.py

@@ -291,6 +310,34 @@ def __init__(self,
            self.model,
            loss=torch.nn.MSELoss())  # will update the loss in successive PR

+    def evaluate_hf(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:


Function to evaluate hf orbitals at the given electron coordinates. (As mentioned in the last PR split the prepare_hf function into smaller ones)

shaipranesh2 · 2023-09-27T15:24:35Z

deepchem/models/torch_models/ferminet.py

+        self.model.loss(up_spin_mo, down_spin_mo, pretrain=True)
+        return np.log(hf_product + np_output**2) + np.log(0.5)
+
+    def pretrain(self,


Function to invoke pretraining

We shouldn't have a custom function for pretraining. Instead follow the convention used elsewhere of having a 'task` parameter and add pretraining/supervised tasks for regular training.

shaipranesh2 · 2023-09-27T15:39:48Z

deepchem/models/torch_models/ferminet.py

@@ -313,3 +360,59 @@ def prepare_hf_solution(self):
        self.mol.build(parse_arg=False)
        self.mf = pyscf.scf.UHF(self.mol)
        _ = self.mf.kernel()
+
+    def random_walk(self, x: np.ndarray) -> np.ndarray:


This function gets called at every step of electron sampling, and it returns the probability of the sampled electrons - log of the wavefunction from the model and the HF solution (this is only for pretraining).

rbharath

@shaipranesh2 We shouldn't have a custom pretrain function but rather should have a task field that we can set. See InfoGraph or other models for examples.

We should also expect to see 22 tests passing but I see 21. Are there any new failures?

shaipranesh2 · 2023-09-29T14:50:13Z

@shaipranesh2 We shouldn't have a custom pretrain function but rather should have a task field that we can set. See InfoGraph or other models for examples.

We should also expect to see 22 tests passing but I see 21. Are there any new failures?

I can see the docs failing in this PR and some others. While it used to passing before (the changes are not due to this PR)

rbharath · 2023-10-02T16:37:45Z

deepchem/models/torch_models/ferminet.py

+             pretrain: List[bool] = [True]):
+        """
+        Implements the loss function for both pretraining and the actual training parts.
+


In a follow up PR, please add more details on the loss. You should be able to list the actual latex for the formulas as well

rbharath · 2023-10-02T16:39:01Z

deepchem/models/torch_models/ferminet.py

@@ -280,7 +311,8 @@ def __init__(self,
            batch_no=self.batch_no,
            central_value=self.nucleon_pos,
            seed=self.seed,
-            f=lambda x: test_f(x),  # Will be replaced in successive PR
+            f=lambda x: self.random_walk(x
+                                        ),  # Will be replaced in successive PR


Make sure to remove the comment in your next update

rbharath

LGTM

shaipranesh2 and others added 14 commits September 3, 2023 19:30

adding hartree fock full support

d4cf0d4

completining ferminet pretrain backward and forward

f3ecd09

minor corrections

c584568

minor mypy fix

996d189

minor mypy fix

59821a7

before rebase

3bd274d

rebased

8006a91

rebased

3a81881

Merge branch 'ferminet_hf' of https://github.com/shaipranesh2/deepchem …

3c88ead

…into ferminet_hf

mypy fix

ee92a6d

fixed conflicts

0b64a1f

mypy fix

e3429c2

small fixes

157af88

small fixes

360f84c

shaipranesh2 force-pushed the ferminet_hf branch from 59821a7 to 3a81881 Compare September 27, 2023 11:02

shaipranesh2 commented Sep 27, 2023

View reviewed changes

rbharath reviewed Sep 27, 2023

View reviewed changes

shaipranesh2 added 7 commits September 28, 2023 19:45

refactoring for tasks

382223f

minor fix

c329f21

minor fix

15b3404

mypy fix

7690834

mypy fix

1a2dd18

mypy fix

92fab49

final changes

c000e3b

rbharath reviewed Oct 2, 2023

View reviewed changes

rbharath approved these changes Oct 2, 2023

View reviewed changes

rbharath merged commit 9183618 into deepchem:master Oct 2, 2023
23 of 33 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

adding complete forward and backward pass for ferminet pretraining #3553

adding complete forward and backward pass for ferminet pretraining #3553

shaipranesh2 commented Sep 3, 2023

shaipranesh2 Sep 27, 2023 •

edited

shaipranesh2 Sep 27, 2023

shaipranesh2 Sep 27, 2023 •

edited

shaipranesh2 Sep 27, 2023

shaipranesh2 Sep 27, 2023 •

edited

shaipranesh2 Sep 27, 2023

rbharath Sep 27, 2023

shaipranesh2 Sep 29, 2023

shaipranesh2 Sep 27, 2023

rbharath left a comment

shaipranesh2 commented Sep 29, 2023 •

edited

rbharath Oct 2, 2023

rbharath Oct 2, 2023

rbharath left a comment

adding complete forward and backward pass for ferminet pretraining #3553

adding complete forward and backward pass for ferminet pretraining #3553

Conversation

shaipranesh2 commented Sep 3, 2023

Description

Type of change

Checklist

shaipranesh2 Sep 27, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shaipranesh2 Sep 27, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shaipranesh2 Sep 27, 2023 • edited

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

rbharath left a comment

Choose a reason for hiding this comment

shaipranesh2 commented Sep 29, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rbharath left a comment

Choose a reason for hiding this comment

shaipranesh2 Sep 27, 2023 •

edited

shaipranesh2 Sep 27, 2023 •

edited

shaipranesh2 Sep 27, 2023 •

edited

shaipranesh2 commented Sep 29, 2023 •

edited