diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 27b6a830a4..d89784e9cf 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -1,22 +1,33 @@ +import os import logging from typing import Callable, Tuple import numpy as np +from deepmd.env import tf from deepmd.utils.errors import OutOfMemoryError + +log = logging.getLogger(__name__) + + class AutoBatchSize: """This class allows DeePMD-kit to automatically decide the maximum batch size that will not cause an OOM error. Notes ----- - We assume all OOM error will raise :class:`OutOfMemoryError`. + 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 :class:`OutOfMemoryError`. Parameters ---------- initial_batch_size : int, default: 1024 - initial batch size (number of total atoms) + initial batch size (number of total atoms) when DP_INFER_BATCH_SIZE + is not set factor : float, default: 2. increased factor @@ -33,8 +44,23 @@ def __init__(self, initial_batch_size: int = 1024, factor: float = 2.) -> None: # See also PyTorchLightning/pytorch-lightning#1638 # TODO: discuss a proper initial batch size self.current_batch_size = initial_batch_size - self.maximum_working_batch_size = 0 - self.minimal_not_working_batch_size = 2**31 + DP_INFER_BATCH_SIZE = int(os.environ.get('DP_INFER_BATCH_SIZE', 0)) + if DP_INFER_BATCH_SIZE > 0: + self.current_batch_size = DP_INFER_BATCH_SIZE + self.maximum_working_batch_size = DP_INFER_BATCH_SIZE + self.minimal_not_working_batch_size = self.maximum_working_batch_size + 1 + else: + self.maximum_working_batch_size = initial_batch_size + if tf.test.is_gpu_available(): + self.minimal_not_working_batch_size = 2**31 + else: + self.minimal_not_working_batch_size = self.maximum_working_batch_size + 1 + log.warning( + "You can use the environment variable DP_INFER_BATCH_SIZE to" + "control the inference batch size (nframes * natoms). " + "The default value is %d." % initial_batch_size + ) + self.factor = factor def execute(self, callable: Callable, start_index: int, natoms: int) -> Tuple[int, tuple]: @@ -86,7 +112,7 @@ def execute(self, callable: Callable, start_index: int, natoms: int) -> Tuple[in def _adjust_batch_size(self, factor: float): old_batch_size = self.current_batch_size self.current_batch_size = int(self.current_batch_size * factor) - logging.info("Adjust batch size from %d to %d" % (old_batch_size, self.current_batch_size)) + log.info("Adjust batch size from %d to %d" % (old_batch_size, self.current_batch_size)) def execute_all(self, callable: Callable, total_size: int, natoms: int, *args, **kwargs) -> Tuple[np.ndarray]: """Excuate a method with all given data. diff --git a/source/tests/test_auto_batch_size.py b/source/tests/test_auto_batch_size.py index f8aa0b8e60..2790c7d4d5 100644 --- a/source/tests/test_auto_batch_size.py +++ b/source/tests/test_auto_batch_size.py @@ -1,4 +1,5 @@ import unittest +import os import numpy as np @@ -11,7 +12,9 @@ def oom(self, batch_size, start_index): raise OutOfMemoryError return batch_size, np.zeros((batch_size, 2)) - def test_execute_oom(self): + @unittest.mock.patch('tensorflow.compat.v1.test.is_gpu_available') + def test_execute_oom_gpu(self, mock_is_gpu_available): + mock_is_gpu_available.return_value = True # initial batch size 256 = 128 * 2 auto_batch_size = AutoBatchSize(256, 2.) # no error - 128 @@ -34,7 +37,48 @@ def test_execute_oom(self): nb, result = auto_batch_size.execute(self.oom, 1, 2) self.assertEqual(nb, 256) self.assertEqual(result.shape, (256, 2)) - + + @unittest.mock.patch('tensorflow.compat.v1.test.is_gpu_available') + def test_execute_oom_cpu(self, mock_is_gpu_available): + mock_is_gpu_available.return_value = False + # initial batch size 256 = 128 * 2, nb is always 128 + auto_batch_size = AutoBatchSize(256, 2.) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + + @unittest.mock.patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "256"}, clear=True) + def test_execute_oom_environment_variables(self): + # DP_INFER_BATCH_SIZE = 256 = 128 * 2, nb is always 128 + auto_batch_size = AutoBatchSize(999, 2.) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + nb, result = auto_batch_size.execute(self.oom, 1, 2) + self.assertEqual(nb, 128) + self.assertEqual(result.shape, (128, 2)) + def test_execute_all(self): dd1 = np.zeros((10000, 2, 1)) auto_batch_size = AutoBatchSize(256, 2.)