[Fix] Fix scale_lr in SingleDeviceStrategy (#1428)

open-mmlab · Nov 20, 2023 · fd5d062 · fd5d062
1 parent 5a90805
commit fd5d062
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/mmengine/_strategy/base.py b/mmengine/_strategy/base.py
@@ -707,7 +707,8 @@ def _is_built(schedulers):
                 return any(_is_built(s) for s in schedulers)
             return isinstance(schedulers, _ParamScheduler)
 
-        if _is_built(self.param_schedulers):
+        if hasattr(self, 'param_schedulers') and _is_built(
+                self.param_schedulers):
             raise RuntimeError('`scale_lr` should be called before building '
                                'ParamScheduler because ParamScheduler will '
                                'store initial lr from optimizer wrappers')

diff --git a/mmengine/_strategy/single_device.py b/mmengine/_strategy/single_device.py
@@ -63,12 +63,6 @@ def prepare(
 
         if optim_wrapper is not None:
             self.optim_wrapper = self.build_optim_wrapper(optim_wrapper, model)
-
-        if param_scheduler is not None:
-            self.param_schedulers = self.build_param_scheduler(
-                param_scheduler, self.optim_wrapper)
-
-        if optim_wrapper is not None:
             self._scale_lr()
 
             accumulative_counts = getattr(self.optim_wrapper,
@@ -82,6 +76,11 @@ def prepare(
 
                 self.optim_wrapper.initialize_count_status(  # type: ignore
                     self.model, 0, self.dispatch_kwargs['max_iters'])
+
+        if param_scheduler is not None:
+            self.param_schedulers = self.build_param_scheduler(
+                param_scheduler, self.optim_wrapper)
+
         self._prepared = True
         return self._prepared_components()
 

diff --git a/mmengine/runner/_flexible_runner.py b/mmengine/runner/_flexible_runner.py
@@ -1176,7 +1176,8 @@ def train(self) -> nn.Module:
             epoch_length=len(self.train_dataloader),
             max_epochs=self.max_epochs,
             max_iters=self.max_iters,
-        )
+            train_micro_batch_size_per_gpu=_get_batch_size(
+                self.train_dataloader))  # type: ignore
 
         self.strategy.prepare(
             self.model,