fix export_model to support dygraph

tools/export_model.py

@@ -15,7 +15,10 @@
 import argparse
 
 from ppcls.modeling import architectures
-import paddle.fluid as fluid
+from ppcls.utils.save_load import load_dygraph_pretrain
+import paddle
+import paddle.nn.functional as F
+from paddle.jit import to_static
 
 
 def parse_args():
@@ -24,54 +27,38 @@ def parse_args():
     parser.add_argument("-p", "--pretrained_model", type=str)
     parser.add_argument("-o", "--output_path", type=str)
     parser.add_argument("--class_dim", type=int, default=1000)
-    parser.add_argument("--img_size", type=int, default=224)
+    # parser.add_argument("--img_size", type=int, default=224)
 
     return parser.parse_args()
 
 
-def create_input(img_size=224):
-    image = fluid.data(
-        name='image', shape=[None, 3, img_size, img_size], dtype='float32')
-    return image
+class Net(paddle.nn.Layer):
+    def __init__(self, net, to_static, class_dim):
+        super(Net, self).__init__()
+        self.pre_net = net(class_dim=class_dim)
+        self.to_static = to_static
 
-
-def create_model(args, model, input, class_dim=1000):
-    if args.model == "GoogLeNet":
-        out, _, _ = model.net(input=input, class_dim=class_dim)
-    else:
-        out = model.net(input=input, class_dim=class_dim)
-        out = fluid.layers.softmax(out)
-    return out
+    # 请根据实际需求修改shape
+    @to_static(input_spec=[
+        paddle.static.InputSpec(
+            shape=[None, 3, 224, 224], dtype='float32')
+    ])
+    def forward(self, inputs):
+        x = self.pre_net(inputs)
+        x = F.softmax(x)
+        return x
 
 
 def main():
     args = parse_args()
 
-    model = architectures.__dict__[args.model]()
-
-    place = fluid.CPUPlace()
-    exe = fluid.Executor(place)
-
-    startup_prog = fluid.Program()
-    infer_prog = fluid.Program()
-
-    with fluid.program_guard(infer_prog, startup_prog):
-        with fluid.unique_name.guard():
-            image = create_input(args.img_size)
-            out = create_model(args, model, image, class_dim=args.class_dim)
-
-    infer_prog = infer_prog.clone(for_test=True)
-    fluid.load(
-        program=infer_prog, model_path=args.pretrained_model, executor=exe)
+    paddle.disable_static()
+    net = architectures.__dict__[args.model]
 
-    fluid.io.save_inference_model(
-        dirname=args.output_path,
-        feeded_var_names=[image.name],
-        main_program=infer_prog,
-        target_vars=out,
-        executor=exe,
-        model_filename='model',
-        params_filename='params')
+    model = Net(net, to_static, args.class_dim)
+    para_state_dict = paddle.io.load_program_state(args.pretrained_model)
+    load_dygraph_pretrain(model, args.pretrained_model, True)
+    paddle.jit.save(model, args.output_path)
 
 
 if __name__ == "__main__":