Merge the preprocess step into the easydeploy onnx/tensorrt model

projects/easydeploy/tools/export_onnx.py

@@ -8,9 +8,10 @@
 import onnx
 import torch
 from mmdet.apis import init_detector
-from mmengine.config import ConfigDict
+from mmengine.config import Config, ConfigDict
 from mmengine.logging import print_log
 from mmengine.utils.path import mkdir_or_exist
+from torch.nn.parameter import Parameter
 
 # Add MMYOLO ROOT to sys.path
 sys.path.append(str(Path(__file__).resolve().parents[3]))
@@ -23,6 +24,32 @@
 warnings.filterwarnings(action='ignore', category=ResourceWarning)
 
 
+def preprocess(config: Config, model: torch.nn.Module):
+    data_preprocess = config.get('model', {}).get('data_preprocessor', {})
+    mean = data_preprocess.get('mean', [0., 0., 0.])
+    std = data_preprocess.get('std', [1., 1., 1.])
+    mean_value = torch.tensor(mean, dtype=torch.float32).reshape(1, 3, 1, 1)
+    std_value = torch.tensor(std, dtype=torch.float32).reshape(1, 3, 1, 1)
+    device = next(model.parameters()).device
+
+    class PreProcess(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+            self.mean_value = Parameter(mean_value, requires_grad=False)
+            self.std_value = Parameter(std_value, requires_grad=False)
+            self.core_model = model
+
+        def forward(self, x: torch.Tensor):
+            assert x.ndim == 4
+            x = x.float()
+            y = (x - self.mean_value) / self.std_value
+            y = self.core_model(y)
+            return y
+
+    return PreProcess().to(device).eval()
+
+
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('config', help='Config file')
@@ -110,6 +137,11 @@ def main():
         baseModel=baseModel, backend=backend, postprocess_cfg=postprocess_cfg)
     deploy_model.eval()
 
+    # embed the preprocess into the model
+    cfg = Config.fromfile(args.config)
+    deploy_model = preprocess(cfg, deploy_model)
+    deploy_model.eval()
+
     fake_input = torch.randn(args.batch_size, 3,
                              *args.img_size).to(args.device)
     # dry run

projects/easydeploy/tools/image-demo.py

@@ -33,27 +33,6 @@ def parse_args():
     return args
 
 
-def preprocess(config):
-    data_preprocess = config.get('model', {}).get('data_preprocessor', {})
-    mean = data_preprocess.get('mean', [0., 0., 0.])
-    std = data_preprocess.get('std', [1., 1., 1.])
-    mean = torch.tensor(mean, dtype=torch.float32).reshape(1, 3, 1, 1)
-    std = torch.tensor(std, dtype=torch.float32).reshape(1, 3, 1, 1)
-
-    class PreProcess(torch.nn.Module):
-
-        def __init__(self):
-            super().__init__()
-
-        def forward(self, x):
-            x = x[None].float()
-            x -= mean.to(x.device)
-            x /= std.to(x.device)
-            return x
-
-    return PreProcess().eval()
-
-
 def main():
     args = parse_args()
 
@@ -80,8 +59,6 @@ def main():
     test_pipeline[0] = ConfigDict({'type': 'mmdet.LoadImageFromNDArray'})
     test_pipeline = Compose(test_pipeline)
 
-    pre_pipeline = preprocess(cfg)
-
     if not args.show:
         path.mkdir_or_exist(args.out_dir)
 
@@ -102,7 +79,7 @@ def main():
             device=args.device)
         scale_factor = samples.get('scale_factor', [1., 1])
         scale_factor = torch.asarray(scale_factor * 2, device=args.device)
-        data = pre_pipeline(data).to(args.device)
+        data = data[None].float().to(args.device)
 
         result = model(data)
         if source_type['is_dir']: