Fix bug that auto adaptive batch size raises an error if CUDA isn't a…

…vailable (#2410) * implement unit test * refine log * Update src/otx/algorithms/common/adapters/mmcv/utils/automatic_bs.py --------- Co-authored-by: Sungman Cho <[email protected]>
openvinotoolkit · Aug 1, 2023 · dcf1081 · dcf1081
1 parent 5c13cff
commit dcf1081
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/src/otx/algorithms/common/adapters/mmcv/utils/automatic_bs.py b/src/otx/algorithms/common/adapters/mmcv/utils/automatic_bs.py
@@ -8,6 +8,7 @@
 from typing import Callable, Dict, List
 
 import numpy as np
+from torch.cuda import is_available as cuda_available
 
 from otx.algorithms.common.adapters.torch.utils import BsSearchAlgo
 from otx.algorithms.common.utils.logger import get_logger
@@ -53,6 +54,10 @@ def adapt_batch_size(train_func: Callable, cfg, datasets: List, validate: bool =
         not_increase (bool) : Whether adapting batch size to larger value than default value or not.
     """
 
+    if not cuda_available():
+        logger.warning("Skip Auto-adaptive batch size: CUDA should be available, but it isn't.")
+        return
+
     def train_func_single_iter(batch_size):
         copied_cfg = deepcopy(cfg)
         _set_batch_size(copied_cfg, batch_size)

diff --git a/tests/unit/algorithms/common/adapters/mmcv/utils/test_automatic_bs.py b/tests/unit/algorithms/common/adapters/mmcv/utils/test_automatic_bs.py
@@ -109,6 +109,19 @@ def test_adapt_batch_size(
     assert len(mock_train_func.call_args_list[0].kwargs["cfg"].custom_hooks) == 1
 
 
+def test_adapt_batch_size_no_gpu(mocker, common_cfg, mock_dataset):
+    # prepare
+    mock_train_func = mocker.MagicMock()
+    mock_config = set_mock_cfg_not_action(common_cfg)
+    mocker.patch.object(automatic_bs, "cuda_available", return_value=False)
+
+    # execute
+    adapt_batch_size(mock_train_func, mock_config, mock_dataset, False, True)
+
+    # check train function ins't called.
+    mock_train_func.assert_not_called()
+
+
 class TestSubDataset:
     @pytest.fixture(autouse=True)
     def set_up(self, mocker):