Merge pull request meituan#800 from triple-Mu/triplemu/fix-upsample

Fix upsample bug
ZHEQIUSHUI · Apr 28, 2023 · 041f547 · 041f547
2 parents bf14e20 + 47cbeb1
commit 041f547
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/deploy/ONNX/export_onnx.py b/deploy/ONNX/export_onnx.py
@@ -51,6 +51,8 @@
     for layer in model.modules():
         if isinstance(layer, RepVGGBlock):
             layer.switch_to_deploy()
+        elif isinstance(layer, nn.Upsample) and not hasattr(layer, 'recompute_scale_factor'):
+            layer.recompute_scale_factor = None  # torch 1.11.0 compatibility
     # Input
     img = torch.zeros(args.batch_size, 3, *args.img_size).to(device)  # image size(1,3,320,192) iDetection
 
@@ -60,7 +62,7 @@
     model.eval()
     for k, m in model.named_modules():
         if isinstance(m, ConvModule):  # assign export-friendly activations
-            if isinstance(m.act, nn.SiLU):
+            if hasattr(m, 'act') and isinstance(m.act, nn.SiLU):
                 m.act = SiLU()
         elif isinstance(m, Detect):
             m.inplace = args.inplace

diff --git a/deploy/OpenVINO/export_openvino.py b/deploy/OpenVINO/export_openvino.py
@@ -42,6 +42,8 @@
     for layer in model.modules():
         if isinstance(layer, RepVGGBlock):
             layer.switch_to_deploy()
+        elif isinstance(layer, nn.Upsample) and not hasattr(layer, 'recompute_scale_factor'):
+            layer.recompute_scale_factor = None  # torch 1.11.0 compatibility
 
     # Input
     img = torch.zeros(args.batch_size, 3, *args.img_size).to(device)  # image size(1,3,320,192) iDetection
@@ -51,8 +53,8 @@
         img, model = img.half(), model.half()  # to FP16
     model.eval()
     for k, m in model.named_modules():
-        if isinstance(m, Conv):  # assign export-friendly activations
-            if isinstance(m.act, nn.SiLU):
+        if isinstance(m, ConvModule):  # assign export-friendly activations
+            if hasattr(m, 'act') and isinstance(m.act, nn.SiLU):
                 m.act = SiLU()
         elif isinstance(m, Detect):
             m.inplace = args.inplace

diff --git a/yolov6/core/evaler.py b/yolov6/core/evaler.py
@@ -64,15 +64,17 @@ def init_model(self, model, weights, task):
                 download_ckpt(weights)
             model = load_checkpoint(weights, map_location=self.device)
             self.stride = int(model.stride.max())
-            if self.device.type != 'cpu':
-                model(torch.zeros(1, 3, self.img_size, self.img_size).to(self.device).type_as(next(model.parameters())))
             # switch to deploy
             from yolov6.layers.common import RepVGGBlock
             for layer in model.modules():
                 if isinstance(layer, RepVGGBlock):
                     layer.switch_to_deploy()
+                elif isinstance(layer, torch.nn.Upsample) and not hasattr(layer, 'recompute_scale_factor'):
+                    layer.recompute_scale_factor = None  # torch 1.11.0 compatibility
             LOGGER.info("Switch model to deploy modality.")
             LOGGER.info("Model Summary: {}".format(get_model_info(model, self.img_size)))
+        if self.device.type != 'cpu':
+            model(torch.zeros(1, 3, self.img_size, self.img_size).to(self.device).type_as(next(model.parameters())))
         model.half() if self.half else model.float()
         return model