deepmd.utils.batch_size

Module Contents

Classes

AutoBatchSize

This class allows DeePMD-kit to automatically decide the maximum

Attributes

log

deepmd.utils.batch_size.log

class deepmd.utils.batch_size.AutoBatchSize(initial_batch_size: int = 1024, factor: float = 2.0)

Bases: abc.ABC

This class allows DeePMD-kit to automatically decide the maximum batch size that will not cause an OOM error.

Parameters:

initial_batch_sizeint, default: 1024

initial batch size (number of total atoms) when DP_INFER_BATCH_SIZE is not set

factorfloat, default: 2.

increased factor

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.

Attributes:

current_batch_sizeint

current batch size (number of total atoms)

maximum_working_batch_sizeint

maximum working batch size

minimal_not_working_batch_sizeint

minimal not working batch size

execute(callable: Callable, start_index: int, natoms: int) → Tuple[int, tuple]

Excuate a method with given batch size.

Parameters:

callableCallable

The method should accept the batch size and start_index as parameters, and returns executed batch size and data.

start_indexint

start index

natomsint

natoms

Returns:

int

executed batch size * number of atoms

tuple

result from callable, None if failing to execute

Raises:

OutOfMemoryError

OOM when batch size is 1

_adjust_batch_size(factor: float)

execute_all(callable: Callable, total_size: int, natoms: int, *args, **kwargs) → Tuple[numpy.ndarray]

Excuate a method with all given data.

Parameters:

callableCallable

The method should accept *args and **kwargs as input and return the similiar array.

total_sizeint

Total size

natomsint

The number of atoms

*args

Variable length argument list.

**kwargs

If 2D np.ndarray, assume the first axis is batch; otherwise do nothing.

abstract is_gpu_available() → bool

Check if GPU is available.

Returns:

bool

True if GPU is available

abstract is_oom_error(e: Exception) → bool

Check if the exception is an OOM error.

Parameters:

eException

Exception

Returns:

bool

True if the exception is an OOM error