for export

Author: horcham

deploy/py_infer/src/core/model/model.py

@@ -106,8 +106,8 @@ def warmup(self):
             height, width = hw_list[0]
             warmup_shape = [(*other_shape, height, width)]  # Only single input
 
-        dummy_tensor = [np.random.randn(*shape).astype(dtype) for shape, dtype in zip(warmup_shape, self.input_dtype)]
-        self.model.infer(dummy_tensor)
+        # dummy_tensor = [np.random.randn(*shape).astype(dtype) for shape, dtype in zip(warmup_shape, self.input_dtype)]
+        # self.model.infer(dummy_tensor)
 
     def __del__(self):
         if hasattr(self, "model") and self.model:

mindocr/losses/det_loss.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from math import pi
 from typing import Tuple, Union
 
@@ -10,6 +11,8 @@
 __all__ = ["DBLoss", "PSEDiceLoss", "EASTLoss", "FCELoss"]
 _logger = logging.getLogger(__name__)
 
+OFFLINE_MODE = os.getenv("OFFLINE_MODE", None)
+
 
 class DBLoss(nn.LossBase):
     """
@@ -165,7 +168,13 @@ def construct(self, pred: Tensor, gt: Tensor, mask: Tensor) -> Tensor:
         neg_loss = (loss * negative).view(loss.shape[0], -1)
 
         neg_vals, _ = ops.sort(neg_loss)
-        neg_index = ops.stack((mnp.arange(loss.shape[0]), neg_vals.shape[1] - neg_count), axis=1)
+
+        if OFFLINE_MODE is None:
+            neg_index = ops.stack((mnp.arange(loss.shape[0]), neg_vals.shape[1] - neg_count), axis=1)
+        else:
+            neg_index = ops.stack(
+                (ops.arange(loss.shape[0], dtype=neg_count.dtype), neg_vals.shape[1] - neg_count), axis=1
+            )
         min_neg_score = ops.expand_dims(ops.gather_nd(neg_vals, neg_index), axis=1)
 
         neg_loss_mask = (neg_loss >= min_neg_score).astype(ms.float32)  # filter values less than top k

mindocr/losses/rec_loss.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 
 import mindspore as ms
@@ -6,6 +8,8 @@
 
 __all__ = ["CTCLoss", "AttentionLoss", "VisionLANLoss"]
 
+OFFLINE_MODE = os.getenv("OFFLINE_MODE", None)
+
 
 class CTCLoss(LossBase):
     """
@@ -147,14 +151,21 @@ class AttentionLoss(LossBase):
     def __init__(self, reduction: str = "mean", ignore_index: int = 0) -> None:
         super().__init__()
         # ignore <GO> symbol, assume it is placed at 0th index
-        self.criterion = nn.CrossEntropyLoss(reduction=reduction, ignore_index=ignore_index)
+        if OFFLINE_MODE is None:
+            self.criterion = nn.CrossEntropyLoss(reduction=reduction, ignore_index=ignore_index)
+        else:
+            self.reduction = reduction
+            self.ignore_index = ignore_index
 
     def construct(self, logits: Tensor, labels: Tensor) -> Tensor:
         labels = labels[:, 1:]  # without <GO> symbol
         num_classes = logits.shape[-1]
         logits = ops.reshape(logits, (-1, num_classes))
         labels = ops.reshape(labels, (-1,))
-        return self.criterion(logits, labels)
+        if OFFLINE_MODE is None:
+            return self.criterion(logits, labels)
+        else:
+            return ops.cross_entropy(logits, labels, reduction=self.reduction, ignore_index=self.ignore_index)
 
 
 class SARLoss(LossBase):

mindocr/models/necks/fpn.py

@@ -1,3 +1,4 @@
+import os
 from typing import List, Tuple
 
 from mindspore import Tensor, nn, ops
@@ -7,14 +8,20 @@
 from ..utils.attention_cells import SEModule
 from .asf import AdaptiveScaleFusion
 
+OFFLINE_MODE = os.getenv("OFFLINE_MODE", None)
 
-def _resize_nn(x: Tensor, scale: int = 0, shape: Tuple[int] = None):
-    if scale == 1 or shape == x.shape[2:]:
-        return x
 
-    if scale:
-        shape = (x.shape[2] * scale, x.shape[3] * scale)
-    return ops.ResizeNearestNeighbor(shape)(x)
+if OFFLINE_MODE is None:
+    def _resize_nn(x: Tensor, scale: int = 0, shape: Tuple[int] = None):
+        if scale == 1 or shape == x.shape[2:]:
+            return x
+
+        if scale:
+            shape = (x.shape[2] * scale, x.shape[3] * scale)
+        return ops.ResizeNearestNeighbor(shape)(x)
+else:
+    def _resize_nn(x: Tensor, shape: Tensor):
+        return ops.ResizeNearestNeighborV2()(x, shape)
 
 
 class FPN(nn.Cell):
@@ -64,11 +71,18 @@ def construct(self, features: List[Tensor]) -> Tensor:
         for i, uc_op in enumerate(self.unify_channels):
             features[i] = uc_op(features[i])
 
-        for i in range(2, -1, -1):
-            features[i] += _resize_nn(features[i + 1], shape=features[i].shape[2:])
+        if OFFLINE_MODE is None:
+            for i in range(2, -1, -1):
+                features[i] += _resize_nn(features[i + 1], shape=features[i].shape[2:])
+
+            for i, out in enumerate(self.out):
+                features[i] = _resize_nn(out(features[i]), shape=features[0].shape[2:])
+        else:
+            for i in range(2, -1, -1):
+                features[i] += _resize_nn(features[i + 1], shape=ops.dyn_shape(features[i])[2:])
 
-        for i, out in enumerate(self.out):
-            features[i] = _resize_nn(out(features[i]), shape=features[0].shape[2:])
+            for i, out in enumerate(self.out):
+                features[i] = _resize_nn(out(features[i]), shape=ops.dyn_shape(features[0])[2:])
 
         return self.fuse(features[::-1])  # matching the reverse order of the original work
 

mindocr/models/transforms/tps_spatial_transformer.py

@@ -1,4 +1,5 @@
 import itertools
+import os
 from typing import Optional, Tuple
 
 import numpy as np
@@ -8,6 +9,8 @@
 import mindspore.ops as ops
 from mindspore import Tensor
 
+OFFLINE_MODE = os.getenv("OFFLINE_MODE", None)
+
 
 def grid_sample(input: Tensor, grid: Tensor, canvas: Optional[Tensor] = None) -> Tensor:
     output = ops.grid_sample(input, grid)
@@ -111,15 +114,22 @@ def __init__(
         self.target_coordinate_repr = Tensor(target_coordinate_repr, dtype=ms.float32)
         self.target_control_points = Tensor(target_control_points, dtype=ms.float32)
 
+        if OFFLINE_MODE is not None:
+            self.matmul = ops.BatchMatMul()
+
     def construct(
         self, input: Tensor, source_control_points: Tensor
     ) -> Tuple[Tensor, Tensor]:
         batch_size = ops.shape(source_control_points)[0]
 
         padding_matrix = ops.tile(self.padding_matrix, (batch_size, 1, 1))
         Y = ops.concat([source_control_points, padding_matrix], axis=1)
-        mapping_matrix = ops.matmul(self.inverse_kernel, Y)
-        source_coordinate = ops.matmul(self.target_coordinate_repr, mapping_matrix)
+        if OFFLINE_MODE is None:
+            mapping_matrix = ops.matmul(self.inverse_kernel, Y)
+            source_coordinate = ops.matmul(self.target_coordinate_repr, mapping_matrix)
+        else:
+            mapping_matrix = self.matmul(self.inverse_kernel[None, ...], Y)
+            source_coordinate = self.matmul(self.target_coordinate_repr[None, ...], mapping_matrix)
         grid = ops.reshape(
             source_coordinate,
             (-1, self.target_height, self.target_width, 2),

mindocr/models/utils/attention_cells.py

@@ -1,3 +1,4 @@
+import os
 from typing import Optional, Tuple
 
 import numpy as np
@@ -9,6 +10,8 @@
 
 __all__ = ["MultiHeadAttention", "PositionwiseFeedForward", "PositionalEncoding", "SEModule"]
 
+OFFLINE_MODE = os.getenv("OFFLINE_MODE", None)
+
 
 class MultiHeadAttention(nn.Cell):
     def __init__(
@@ -108,9 +111,14 @@ def __init__(
         self.pe = Tensor(pe, dtype=ms.float32)
 
     def construct(self, input_tensor: Tensor) -> Tensor:
-        input_tensor = (
-            input_tensor + self.pe[:, : input_tensor.shape[1]]
-        )  # pe 1 5000 512
+        if OFFLINE_MODE is None:
+            input_tensor = (
+                input_tensor + self.pe[:, : input_tensor.shape[1]]
+            )  # pe  1 5000 512
+        else:
+            input_tensor = (
+                input_tensor + self.pe[:, : ops.dyn_shape(input_tensor)[1]]
+            )  # pe 1 5000 512
         return self.dropout(input_tensor)