usnistgov / alignn

Atomistic Line Graph Neural Network https://scholar.google.com/citations?user=9Q-tNnwAAAAJ&hl=en
https://jarvis.nist.gov/jalignn/
Other
219 stars 80 forks source link

Make pretrained.py (or similar) usable for ALIGNN FF #150

Open ironhammer269 opened 5 months ago

ironhammer269 commented 5 months ago

As of right now, the only way to perform inference on ALIGNN FF models is to use the run_alignn_ff script.

It is practical if one is interested in performing MD on a single structure, but for prediction of simple properties for multiple structures (eg. unrelaxed energy), it is very inefficient since one has to reload the model for every inference. Also, it is not easily parallelized as the code for calculating the energy is locked in another module after "self.atoms = self.jarvis_atoms.ase_converter()" line 367 in ff.py.

Thus i have been trying to use pretrained.py for an ALIGNN FF model but i am running into the following issue :

(alignn_env) PS C:\Users\haear\Documents\GitHub\alignn\alignn> python .\pretrained.py --model_path "C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed" --file_format poscar --file_path "C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\DFT\test_multi_pass\POSCAR1" Using chk file C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed\best_model.pt Path C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed Config C:\Users\haear\Documents\polytechnique\stage3a\scienti\eval_energy\models\alex_force_fixed\config.json C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\dgl\backend\pytorch\tensor.py:445: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly. To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage() assert input.numel() == input.storage().size(), ( Traceback (most recent call last): File "C:\Users\haear\Documents\GitHub\alignn\alignn\pretrained.py", line 506, in out_data = get_prediction( File "C:\Users\haear\Documents\GitHub\alignn\alignn\pretrained.py", line 344, in get_prediction model([g.to(device), lg.to(device)])['out'] File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, *kwargs) File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl return forward_call(args, **kwargs) File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\alignn\models\alignn_atomwise.py", line 524, in forward / g.ndata["V"][count_node + num_nodes] File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\dgl\view.py", line 80, in getitem return self._graph._get_n_repr(self._ntid, self._nodes)[key] File "C:\Users\haear\anaconda3\envs\alignn_env\lib\site-packages\dgl\frame.py", line 688, in getitem return self._columns[name].data KeyError: 'V'

I guess this is linked to the difference in construction of the graphs between ALIGNN FF/ non FF, but i haven't been able to track down what it is. Also this difference is mainly a decoration since both models should take the same input for inference (ie. a simple POSCAR file).

If anyone has insight on how to fix this, I would be very grateful.