deepmd.pt.utils.auto_batch_size

Module Contents

Classes

AutoBatchSize

Auto batch size.

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

Bases: deepmd.utils.batch_size.AutoBatchSize

Auto batch size.

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

is_gpu_available() → bool

Check if GPU is available.

Returns:

bool

True if GPU is available

is_oom_error(e: Exception) → bool

Check if the exception is an OOM error.

Parameters:

eException

Exception

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

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 or torch.Tensor, assume the first axis is batch; otherwise do nothing.