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