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