Comments (9)
That should be pretty trivial, just add to utils.py
a loop of grad. Might need to do some matrix reconstruction though.
from mace.
I don't think analytical Hessian belongs in the ASE calculator, at least for now - it's useful, but it's not a standard ASE calculator result. It is trivial to do naively, or you can use https://wiki.fysik.dtu.dk/ase/ase/vibrations/modes.html
, or use phonopy
to include symmetry to reduce the number of finite differences.
from mace.
I am aware that there is not too much of an interested in implementing Hessians, but a proper implementation would be very useful. Often one of the selling points of MLIP are the easy to get Hessians and therefore there should be a community interested in that. So was trying to implement hessians for MACE and I was hitting some problems. As recommended from @ilyes319 a loop is a way to implement hessian:
hessian1=[]
for grad_elem in forces.view(-1):
print(grad_elem)
hess_row = torch.autograd.grad(
outputs = [grad_elem],
inputs = [positions],
grad_outputs=torch.ones_like(grad_elem),
retain_graph=True,
create_graph=False,
allow_unused=False,
)[0]
hess_row =hess_row.detach() #this makes it very slow? but needs less memory
if hess_row is None:
hessian1.append(torch.zeros_like(positions))
else:
hessian1.append(hess_row.view(1, -1))
This works very good!, but problem here is, for larger systems it takes very long to iterate through, here time is the limiting factor.
It is also possible to make it vectorized, which is significantly faster but also needs a bit better understanding of the autograd.grad function:
tuple_of_tensors=torch.eye(forces.view(-1).shape[0], forces.view(-1).shape[0],device="cuda")
hessian= torch.autograd.grad(
outputs = [forces.view(-1)],#*45,# forces.view(-1),
inputs = (positions,),
grad_outputs=(tuple_of_tensors,),
retain_graph=True,
create_graph=False,
allow_unused=False,
is_grads_batched=True,
)
Here the is_grads_batched allows a vectorized behavior, but this has the problem of being extremely memory intensive.
I also tried implementations based torch.autograd.functional.hessian
ortorch.autograd.functional.jacobian
or the torch.func.hessian
function with rewritting larger parts of MACE into a function that yields energy and forces respectively.
Nevertheless, non of these implementations is memory efficient and fast at the same time
For systems with more than 100 atoms they are either to slow or to memory intense. Maybe someone here, with a better understanding of the MACE code and pytorch knows a solution for fast and memory efficient Hessians for MACE.
from mace.
Indeed computing analytical Hessian is quite expensive. I think these are mostly all the exact methods from torch. Can you give me an estimate of the sizes of systems you want to look at?
Also can you tell me the GPU you are using and the timing for each of the methods.
from mace.
I am using a A100 with 80Gb RAM
My current testing code looks like this:
from mace.calculators import mace_mp
from ase import build
atoms = build.molecule('H2O')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float32",device='cuda' )
atoms.calc = calc
#print(atoms.get_potential_energy())
#exit()
from ase.build import fcc111
initial = fcc111('Pt', size=(7, 8, 1), vacuum=10.0, orthogonal=True) # size i would like to test = (7,8,6)
initial.calc = calc
e0=initial.get_potential_energy()
print(e0)
For the loop approach this takes ~ 4.5 seconds (could restructure my code then I could give you a more precise answer).
For the vectorized approach this takes ~1.5 seconds.
With the size of the system I am actually working with (for example size = (7,8,6))
loop: ~85 seconds
vectorized: Memory overflow
this is ONLY the hessian calculation by the autograd,grad!
If i use autograd.functional. hessian:
h = torch.autograd.functional.hessian(
self.get_energy,
inputs=data["positions"],
vectorize = False,
#outer_jacobian_strategy= 'reverse-mode'
)
System size = (7,8,1), vecorize =False: ~6.8 seconds
System size = (7,8,1), vecorize =True: ~1.8 seconds
System size = (7,8,6), vecorize =False: ~87 seconds
System size = (7,8,6), vecorize =True: Memory overflow
h = torch.autograd.functional.jacobian(self.get_forces,data["positions"], vectorize=False)
System size = (7,8,1), vecorize =False: ~6.8 seconds
System size = (7,8,1), vecorize =True: ~1.8 seconds
The two torch.func approaches
h=torch.func.hessian(self.get_energy)(data["positions"])#,data["node_attrs"])
h=torch.func.jacrev(torch.func.jacfwd(self.get_energy),chunk_size= 1)(data["positions"])#,data["node_attrs"])
are both very similar in speed ~1.3 seconds but sometimes have already a memory overflow for the system size (7,8,1).
I am overall not the biggest fan of rewritting too much of the code so all options with get_energy or get_forces are rather experimental.
Thanks for the fast answer I hope this stuff helps.
from mace.
One thing I haven't noticed in this discussion is whether you're trying to compute the entire Hessian matrix explicitly, or whether you're using symmetry to take advantage of the fact that all atoms are equivalent in your fcc cell and you only need phonopy
does automatically, e.g.)
from mace.
While my testing system is a fcc cell, I will also need Hessian matrices of disordered systems and therefore I need them explicitly. Ofc an estimation method could be sufficient but the reason why I want to use a MLIP is this idea of cheap and fast automated derivatives.
from mace.
Have you tried batching the rows, so that it does a few at a time, minimizing overhead, but doesn't exceed memory?
from mace.
I indeed tried to do it in batches and it seems to be the best solution for now but it doesn't feel to like a good solution but maybe it is the best for now.
Edit: And it is also still very inefficent
from mace.
Related Issues (20)
- Can we use pretrain model (mace_mp) and retrain with our configurations? HOT 2
- Multi-GPU evaluation for optimisation of a full protein structure HOT 1
- checkpoint file not read in foundation branch. HOT 7
- Error while using foundation branch HOT 8
- Update install instructions for conda-forge HOT 5
- loss function evaluated as zero. HOT 2
- Setting `dispersion=True` in the MACE-MP-0 foundation model omits data from `calc.results` HOT 1
- LAMMPS-MACE with Kokkos: Illegal memory access encountered with trained MACE but works with MP-MACE-0 HOT 51
- multi-gpu lammps issue HOT 6
- Fine-tunning up-to-date MACE-MP-0 model HOT 7
- Don't modify torch default dtype HOT 1
- typo? in multi-gpu branch run_train.py
- repulsion branch DistributedEnvironment exception produces no output HOT 1
- Where are the Similarity Statement Files for Mace-mp-0? HOT 1
- How to repeat the test-result in your paper
- Possibility to not compute energies on `MACE.forward` HOT 1
- run_train.py crashes if swa fails to make progress HOT 2
- RMSE increases when training on multi-gpu HOT 1
- questions about mace documentation HOT 6
- How to add two MACE models in `mace-md`? 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 mace.