未验证 提交 53c5850d 编写于 作者: L littletomatodonkey 提交者: GitHub

Merge pull request #291 from TingquanGao/ppcls-dev

fix export_model to support dygraph
...@@ -15,63 +15,57 @@ ...@@ -15,63 +15,57 @@
import argparse import argparse
from ppcls.modeling import architectures 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(): def parse_args():
def str2bool(v):
return v.lower() in ("true", "t", "1")
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("-m", "--model", type=str) parser.add_argument("-m", "--model", type=str)
parser.add_argument("-p", "--pretrained_model", type=str) parser.add_argument("-p", "--pretrained_model", type=str)
parser.add_argument("-o", "--output_path", type=str) parser.add_argument("-o", "--output_path", type=str)
parser.add_argument("--class_dim", type=int, default=1000) parser.add_argument("--class_dim", type=int, default=1000)
parser.add_argument("--img_size", type=int, default=224) parser.add_argument("--load_static_weights", type=str2bool, default=True)
# parser.add_argument("--img_size", type=int, default=224)
return parser.parse_args() return parser.parse_args()
def create_input(img_size=224): class Net(paddle.nn.Layer):
image = fluid.data( def __init__(self, net, to_static, class_dim):
name='image', shape=[None, 3, img_size, img_size], dtype='float32') super(Net, self).__init__()
return image self.pre_net = net(class_dim=class_dim)
self.to_static = to_static
def create_model(args, model, input, class_dim=1000): # Please modify the 'shape' according to actual needs
if args.model == "GoogLeNet": @to_static(input_spec=[
out, _, _ = model.net(input=input, class_dim=class_dim) paddle.static.InputSpec(
else: shape=[None, 3, 224, 224], dtype='float32')
out = model.net(input=input, class_dim=class_dim) ])
out = fluid.layers.softmax(out) def forward(self, inputs):
return out x = self.pre_net(inputs)
x = F.softmax(x)
return x
def main(): def main():
args = parse_args() args = parse_args()
model = architectures.__dict__[args.model]() paddle.disable_static()
net = 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) model = Net(net, to_static, args.class_dim)
fluid.load(
program=infer_prog, model_path=args.pretrained_model, executor=exe)
fluid.io.save_inference_model( load_dygraph_pretrain(
dirname=args.output_path, model.pre_net,
feeded_var_names=[image.name], path=args.pretrained_model,
main_program=infer_prog, load_static_weights=args.load_static_weights)
target_vars=out, paddle.jit.save(model, args.output_path)
executor=exe,
model_filename='model',
params_filename='params')
if __name__ == "__main__": if __name__ == "__main__":
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册