do not apply automatic batch size on CPUs

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. 

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.)