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).
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.
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.
#####################################################