Fix error during loss reduce in callback

mindocr/utils/callbacks.py

@@ -9,7 +9,7 @@
 from .checkpoint import CheckpointManager
 from .evaluator import Evaluator
 from .logger import Logger
-from .misc import AverageMeter, fetch_optimizer_lr
+from .misc import AllReduce, AverageMeter, fetch_optimizer_lr
 from .recorder import PerfRecorder
 
 __all__ = ["EvalSaveCallback"]
@@ -95,8 +95,8 @@ def __init__(
 
         self._loss_avg_meter = AverageMeter()
 
-        self._reduce_sum = ms.ops.AllReduce()
         self._device_num = device_num
+        self._reduce = AllReduce(device_num=self._device_num)
         # lamda expression is not supported in jit
         self._loss_reduce = self._reduce if device_num is not None else lambda x: x
 
@@ -110,9 +110,6 @@ def __init__(
             )
         self.start_epoch = start_epoch
 
-    def _reduce(self, x):
-        return self._reduce_sum(x) / self._device_num  # average value across all devices
-
     def on_train_step_end(self, run_context):
         """
         Print training loss at the end of step.

mindocr/utils/misc.py

@@ -1,6 +1,10 @@
+from typing import Optional
+
 from packaging import version
 
 import mindspore as ms
+import mindspore.nn as nn
+import mindspore.ops as ops
 from mindspore import Tensor
 from mindspore.ops.primitive import constexpr
 
@@ -39,6 +43,28 @@ def fetch_optimizer_lr(opt):
     return lr
 
 
+class AllReduce(nn.Cell):
+    def __init__(self, reduce: str = "mean", device_num: Optional[int] = None) -> None:
+        super().__init__()
+        self.average = reduce == "mean"
+
+        if device_num is None:
+            self.device_num = 1
+        else:
+            self.device_num = device_num
+
+        self.all_reduce = ops.AllReduce()
+
+    def construct(self, x: Tensor) -> Tensor:
+        dtype = x.dtype
+        x = ops.cast(x, ms.float32)
+        x = self.all_reduce(x)
+        if self.average:
+            x = x / self.device_num
+        x = ops.cast(x, dtype)
+        return x
+
+
 @constexpr
 def is_ms_version_2():
     """This check can be applied in `nn.Cell.construct` method, to