From f04b52343ff9b70c008aa5605b886e8a292e887f Mon Sep 17 00:00:00 2001
From: gaotingquan <gaotingquan@baidu.com>
Date: Fri, 13 May 2022 15:41:08 +0000
Subject: [PATCH] fix: unable to export rep net

---
 ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py | 8 +++-----
 ppcls/arch/backbone/model_zoo/repvgg.py             | 9 ++-------
 ppcls/engine/engine.py                              | 6 ++++++
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py b/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py
index 3ce03a9c9..459d84275 100644
--- a/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py
+++ b/ppcls/arch/backbone/legendary_models/pp_lcnet_v2.py
@@ -188,7 +188,7 @@ class RepDepthwiseSeparable(TheseusLayer):
     def forward(self, x):
         if self.use_rep:
             input_x = x
-            if not self.training:
+            if self.is_repped:
                 x = self.act(self.dw_conv(x))
             else:
                 y = self.dw_conv_list[0](x)
@@ -209,14 +209,12 @@ class RepDepthwiseSeparable(TheseusLayer):
             x = x + input_x
         return x
 
-    def eval(self):
+    def rep(self):
         if self.use_rep:
+            self.is_repped = True
             kernel, bias = self._get_equivalent_kernel_bias()
             self.dw_conv.weight.set_value(kernel)
             self.dw_conv.bias.set_value(bias)
-        self.training = False
-        for layer in self.sublayers():
-            layer.eval()
 
     def _get_equivalent_kernel_bias(self):
         kernel_sum = 0
diff --git a/ppcls/arch/backbone/model_zoo/repvgg.py b/ppcls/arch/backbone/model_zoo/repvgg.py
index 8ff662a7f..12f65549f 100644
--- a/ppcls/arch/backbone/model_zoo/repvgg.py
+++ b/ppcls/arch/backbone/model_zoo/repvgg.py
@@ -124,13 +124,7 @@ class RepVGGBlock(nn.Layer):
             groups=groups)
 
     def forward(self, inputs):
-        if not self.training and not self.is_repped:
-            self.rep()
-            self.is_repped = True
-        if self.training and self.is_repped:
-            self.is_repped = False
-
-        if not self.training:
+        if self.is_repped:
             return self.nonlinearity(self.rbr_reparam(inputs))
 
         if self.rbr_identity is None:
@@ -154,6 +148,7 @@ class RepVGGBlock(nn.Layer):
         kernel, bias = self.get_equivalent_kernel_bias()
         self.rbr_reparam.weight.set_value(kernel)
         self.rbr_reparam.bias.set_value(bias)
+        self.is_repped = True
 
     def get_equivalent_kernel_bias(self):
         kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
diff --git a/ppcls/engine/engine.py b/ppcls/engine/engine.py
index 5b5c4da8a..675e92a6d 100644
--- a/ppcls/engine/engine.py
+++ b/ppcls/engine/engine.py
@@ -452,6 +452,12 @@ class Engine(object):
                                   self.config["Global"]["pretrained_model"])
 
         model.eval()
+
+        # for rep nets
+        for layer in self.model.sublayers():
+            if hasattr(layer, "rep"):
+                layer.rep()
+
         save_path = os.path.join(self.config["Global"]["save_inference_dir"],
                                  "inference")
         if model.quanter: