Bugfix crnn

mindocr/models/builder.py

@@ -7,7 +7,7 @@
 
 from ._registry import is_model, list_models, model_entrypoint
 from .base_model import BaseModel
-from .utils import load_model
+from .utils import load_model, set_amp_attr
 
 __all__ = ["build_model"]
 
@@ -74,5 +74,6 @@ def build_model(name_or_config: Union[str, dict], **kwargs):
 
     if "amp_level" in kwargs:
         auto_mixed_precision(network, amp_level=kwargs["amp_level"])
+        set_amp_attr(network, kwargs["amp_level"])
 
     return network

mindocr/models/necks/rnn.py

@@ -2,7 +2,9 @@
 
 import numpy as np
 
-from mindspore import Tensor, nn, ops
+import mindspore.ops.functional as F
+from mindspore import Tensor, nn, ops, version
+from mindspore.common import dtype
 
 __all__ = ['RNNEncoder']
 
@@ -37,6 +39,11 @@ def __init__(self, in_channels: int, hidden_size: int = 512,  batch_size: Option
                                    has_bias=True,
                                    dropout=0.,
                                    bidirectional=True)
+        self.encoder_cast_to_fp16 = False
+        if version.__version__ >= "2.3":
+            # Adapted to MindSpore r2.3, nn.LSTM has bugs when input is FP32.
+            self.seq_encoder.to_float(dtype.float16)
+            self.encoder_cast_to_fp16 = True
 
         self.hx = None
         if batch_size is not None:
@@ -49,9 +56,15 @@ def construct(self, features: List[Tensor]) -> Tensor:
         x = ops.squeeze(x, axis=2)  # [N, C, W]
         x = ops.transpose(x, (2, 0, 1))  # [W, N, C]
 
+        if self.encoder_cast_to_fp16 and self._amp_level == "O0":
+            x = F.cast(x, dtype.float16)
+
         if self.hx is None:
             x, _ = self.seq_encoder(x)
         else:
             x, _ = self.seq_encoder(x, self.hx)
 
-        return x
+        if self.encoder_cast_to_fp16 and self._amp_level == "O0":
+            return F.cast(x, dtype.float32)
+        else:
+            return x

mindocr/models/utils/__init__.py

@@ -1,3 +1,3 @@
 from .attention_cells import *
-from .load_model import load_model
+from .load_model import load_model, set_amp_attr
 from .rnn_cells import GRUCell

mindocr/models/utils/load_model.py

@@ -2,11 +2,11 @@
 import os
 from typing import Callable, Dict, Optional
 
-from mindspore import load_checkpoint, load_param_into_net
+from mindspore import load_checkpoint, load_param_into_net, nn
 
 from ..backbones.mindcv_models.utils import auto_map, download_pretrained
 
-__all__ = ["load_model", "drop_inconsistent_shape_parameters"]
+__all__ = ["load_model", "drop_inconsistent_shape_parameters", "set_amp_attr"]
 _logger = logging.getLogger(__name__)
 
 
@@ -78,3 +78,9 @@ def load_model(
         f"Finish loading model checkoint from {load_from}. "
         "If no parameter fail-load warning displayed, all checkpoint params have been successfully loaded."
     )
+
+
+def set_amp_attr(network : nn.Cell, amp_level : str):
+    cells = network.name_cells()
+    for name in cells:
+        setattr(network._cells[name], "_amp_level", amp_level)