Single GPU Support¶
Warning The use of GPU and mutli-GPU is under developpement and hasn’t been thoroughly tested yet. Proceed with caution !
Using pytorch as a backend, QMCTorch can leverage GPU cards available on your hardware. You of course must have the CUDA version of pytorch installed (https://pytorch.org/)
Let’s first import everything and create a molecule
[13]:
import torch
from torch import optim
from qmctorch.scf import Molecule
from qmctorch.wavefunction import SlaterJastrow
from qmctorch.sampler import Metropolis
from qmctorch.utils import (plot_energy, plot_data)
mol = Molecule(atom='H 0. 0. 0; H 0. 0. 1.', unit='bohr', redo_scf=True)
INFO:QMCTorch|
INFO:QMCTorch| SCF Calculation
INFO:QMCTorch| Removing H2_pyscf_sto-3g.hdf5 and redo SCF calculations
INFO:QMCTorch| Running scf calculation
converged SCF energy = -1.06599946214331
INFO:QMCTorch| Molecule name : H2
INFO:QMCTorch| Number of electrons : 2
INFO:QMCTorch| SCF calculator : pyscf
INFO:QMCTorch| Basis set : sto-3g
INFO:QMCTorch| SCF : HF
INFO:QMCTorch| Number of AOs : 2
INFO:QMCTorch| Number of MOs : 2
INFO:QMCTorch| SCF Energy : -1.066 Hartree
The use of GPU acceleration has been streamlined in QMCTorch, the only modification you need to do on your code is to specify cuda=True
in the declaration of the wave function and sampler, this will automatically port all the necesaary tensors to the GPU and offload all the corresponding operation there.
[9]:
if torch.cuda.is_available():
wf = SlaterJastrow(mol, cuda=True)
sampler = Metropolis(nwalkers=100, nstep=500, step_size=0.25,
nelec=wf.nelec, ndim=wf.ndim,
init=mol.domain('atomic'),
move={'type': 'all-elec', 'proba': 'normal'},
cuda=True)
else:
print('CUDA not available, install torch with cuda support to proceed')
CUDA not available, install torch with cuda support to proceed