deepmd.utils.batch_size#
Attributes#
Classes#
This class allows DeePMD-kit to automatically decide the maximum |
Module Contents#
- class deepmd.utils.batch_size.AutoBatchSize(initial_batch_size: int = 1024, factor: float = 2.0)[source]#
Bases:
abc.ABC
This class allows DeePMD-kit to automatically decide the maximum batch size that will not cause an OOM error.
- Parameters:
- Attributes:
Notes
In some CPU environments, the program may be directly killed when OOM. In this case, by default the batch size will not be increased for CPUs. The environment variable DP_INFER_BATCH_SIZE can be set as the batch size.
In other cases, we assume all OOM error will raise
OutOfMemoryError
.- execute(callable: Callable, start_index: int, natoms: int) tuple[int, tuple] [source]#
Excuate a method with given batch size.
- Parameters:
- Returns:
- Raises:
OutOfMemoryError
OOM when batch size is 1
- execute_all(callable: Callable, total_size: int, natoms: int, *args, **kwargs) tuple[numpy.ndarray] [source]#
Excuate a method with all given data.
This method is compatible with Array API.