CPU & GPU (CUDA), parallelization¶
If the GPU-enabled versions of the backends have been installed, then by default the calculation will try to autodetect if there is a GPU available
to run on, and if so use that. But you can also enforce to run on a specific
Device
.
To limit the number of CPU threads, the NumThreads
keyword can be used if the
backend uses PyTorch as its machine learning framework. Alternatively, you can
set the environment variable OMP_NUM_THREADS.
To use a CUDA-enabled GPU, ensure that a CUDA-enabled version of TensorFlow or PyTorch has been installed (see Installation & Uninstallation).
If the software has problems detecting the GPU, see Debugging installation and available resources.
Device
- Type:
Multiple Choice
- Default value:
- Options:
[, cpu, cuda:0, cuda:1, mps]
- Description:
Device on which to run the calculation (e.g. cpu, cuda:0). If empty, the device can be controlled using environment variables for TensorFlow or PyTorch.
NumThreads
- Type:
String
- Default value:
- GUI name:
Number of threads
- Description:
Number of threads. If not empty, OMP_NUM_THREADS will be set to this number; for PyTorch-engines, torch.set_num_threads() will be called.
Note
Because the calculation runs in a separate process, the number of threads is controlled by the input keyword NumThreads and not by the environment variable NSCM. We recommend setting NSCM=1 when using the MLPotential engine.
Only single-node calculations are currently supported.
SLURM users may need to set SCM_DISABLE_MPI=1.
AMS will report the compute resources it found for machine learning potentials in standard out after applying any restrictions from environment variables as well as NumThreads
and Device
.
GPU with PyTorch backend:
#####################PyTorch setup###################
Number of threads were not limited, using all available CPU cores.
Using GPU: "cuda:0" as found by auto setup.
####################################################
GPU with TensorFlow backend:
###################TensorFlow setup##################
TensorFlow found the following devices:
PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')
PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')
NumThreads was not specified in the MLPotential engine block, so TensorFlow will use all available cores.
#####################################################
CPU only with PyTorch backend:
#####################PyTorch setup###################
Number of threads were not limited, using all available CPU cores.
Using CPU only.
#####################################################
CPU only with TensorFlow backend:
###################TensorFlow setup##################
TensorFlow found the following devices:
PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')
NumThreads was not specified in the MLPotential engine block, so TensorFlow will use all available cores.
#####################################################