From c4912fba3f27b7ad915ceffdb1d8103a3b039f1d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Dec 2022 19:07:51 -0500 Subject: [PATCH 1/4] add an environmental variable to control inference batch size --- deepmd/utils/batch_size.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 27b6a830a4..1aff6e0ded 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -1,3 +1,4 @@ +import os import logging from typing import Callable, Tuple @@ -11,7 +12,11 @@ class AutoBatchSize: 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, the environment variable `DP_INFER_BATCH_SIZE` is used as the + batch size. + + In other cases, we assume all OOM error will raise :class:`OutOfMemoryError`. Parameters ---------- @@ -33,8 +38,13 @@ 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.maximum_working_batch_size = DP_INFER_BATCH_SIZE + self.minimal_not_working_batch_size = DP_INFER_BATCH_SIZE + 1 + else: + self.maximum_working_batch_size = initial_batch_size + self.minimal_not_working_batch_size = 2**31 self.factor = factor def execute(self, callable: Callable, start_index: int, natoms: int) -> Tuple[int, tuple]: From 98b6096cfe4888df9ce843459d15dd97376f10fb Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 30 Nov 2022 08:02:31 -0500 Subject: [PATCH 2/4] do not increase the batch size for CPUs Signed-off-by: Jinzhe Zeng --- deepmd/utils/batch_size.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 1aff6e0ded..e0d23b2cf4 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -4,8 +4,13 @@ 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. @@ -13,15 +18,16 @@ class AutoBatchSize: Notes ----- In some CPU environments, the program may be directly killed when OOM. In - this case, the environment variable `DP_INFER_BATCH_SIZE` is used as the - batch size. + 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 @@ -41,10 +47,19 @@ def __init__(self, initial_batch_size: int = 1024, factor: float = 2.) -> None: DP_INFER_BATCH_SIZE = int(os.environ.get('DP_INFER_BATCH_SIZE', 0)) if DP_INFER_BATCH_SIZE > 0: self.maximum_working_batch_size = DP_INFER_BATCH_SIZE - self.minimal_not_working_batch_size = DP_INFER_BATCH_SIZE + 1 + self.minimal_not_working_batch_size = self.maximum_working_batch_size + 1 else: self.maximum_working_batch_size = initial_batch_size - self.minimal_not_working_batch_size = 2**31 + 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]: @@ -96,7 +111,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. From a17f82f30480b60bc59bc89a2e8a8560614ec68b Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 30 Nov 2022 08:30:11 -0500 Subject: [PATCH 3/4] set current_batch_size to DP_INFER_BATCH_SIZE Signed-off-by: Jinzhe Zeng --- deepmd/utils/batch_size.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index e0d23b2cf4..d89784e9cf 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -46,6 +46,7 @@ def __init__(self, initial_batch_size: int = 1024, factor: float = 2.) -> None: self.current_batch_size = initial_batch_size 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: From b9fa8192f310bca5ef84c25cf56175d7027a6c8a Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 2 Dec 2022 01:51:19 -0500 Subject: [PATCH 4/4] fix tests Signed-off-by: Jinzhe Zeng --- source/tests/test_auto_batch_size.py | 48 ++++++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) 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.)