load_model.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import yaml
  15. import os.path as osp
  16. import six
  17. import copy
  18. from collections import OrderedDict
  19. import paddle.fluid as fluid
  20. from paddle.fluid.framework import Parameter
  21. import paddlex
  22. import paddlex.utils.logging as logging
  23. from paddlex.cv.transforms import build_transforms, build_transforms_v1
  24. def load_model(model_dir, fixed_input_shape=None):
  25. if not osp.exists(osp.join(model_dir, "model.yml")):
  26. raise Exception("There's not model.yml in {}".format(model_dir))
  27. with open(osp.join(model_dir, "model.yml")) as f:
  28. info = yaml.load(f.read(), Loader=yaml.Loader)
  29. if 'status' in info:
  30. status = info['status']
  31. elif 'save_method' in info:
  32. # 兼容老版本PaddleX
  33. status = info['save_method']
  34. if not hasattr(paddlex.cv.models, info['Model']):
  35. raise Exception("There's no attribute {} in paddlex.cv.models".format(
  36. info['Model']))
  37. if 'model_name' in info['_init_params']:
  38. del info['_init_params']['model_name']
  39. model = getattr(paddlex.cv.models, info['Model'])(**info['_init_params'])
  40. model.fixed_input_shape = fixed_input_shape
  41. if '_Attributes' in info:
  42. if 'fixed_input_shape' in info['_Attributes']:
  43. fixed_input_shape = info['_Attributes']['fixed_input_shape']
  44. if fixed_input_shape is not None:
  45. logging.info("Model already has fixed_input_shape with {}".
  46. format(fixed_input_shape))
  47. model.fixed_input_shape = fixed_input_shape
  48. if status == "Normal" or \
  49. status == "Prune" or status == "fluid.save":
  50. startup_prog = fluid.Program()
  51. model.test_prog = fluid.Program()
  52. with fluid.program_guard(model.test_prog, startup_prog):
  53. with fluid.unique_name.guard():
  54. model.test_inputs, model.test_outputs = model.build_net(
  55. mode='test')
  56. model.test_prog = model.test_prog.clone(for_test=True)
  57. model.exe.run(startup_prog)
  58. if status == "Prune":
  59. from .slim.prune import update_program
  60. model.test_prog = update_program(model.test_prog, model_dir,
  61. model.places[0])
  62. import pickle
  63. with open(osp.join(model_dir, 'model.pdparams'), 'rb') as f:
  64. load_dict = pickle.load(f)
  65. fluid.io.set_program_state(model.test_prog, load_dict)
  66. elif status == "Infer" or \
  67. status == "Quant" or status == "fluid.save_inference_model":
  68. [prog, input_names, outputs] = fluid.io.load_inference_model(
  69. model_dir, model.exe, params_filename='__params__')
  70. model.test_prog = prog
  71. test_outputs_info = info['_ModelInputsOutputs']['test_outputs']
  72. model.test_inputs = OrderedDict()
  73. model.test_outputs = OrderedDict()
  74. for name in input_names:
  75. model.test_inputs[name] = model.test_prog.global_block().var(name)
  76. for i, out in enumerate(outputs):
  77. var_desc = test_outputs_info[i]
  78. model.test_outputs[var_desc[0]] = out
  79. if 'Transforms' in info:
  80. transforms_mode = info.get('TransformsMode', 'RGB')
  81. # 固定模型的输入shape
  82. fix_input_shape(info, fixed_input_shape=fixed_input_shape)
  83. if transforms_mode == 'RGB':
  84. to_rgb = True
  85. else:
  86. to_rgb = False
  87. if 'BatchTransforms' in info:
  88. # 兼容老版本PaddleX模型
  89. model.test_transforms = build_transforms_v1(
  90. model.model_type, info['Transforms'], info['BatchTransforms'])
  91. model.eval_transforms = copy.deepcopy(model.test_transforms)
  92. else:
  93. model.test_transforms = build_transforms(
  94. model.model_type, info['Transforms'], to_rgb)
  95. model.eval_transforms = copy.deepcopy(model.test_transforms)
  96. if '_Attributes' in info:
  97. for k, v in info['_Attributes'].items():
  98. if k in model.__dict__:
  99. model.__dict__[k] = v
  100. logging.info("Model[{}] loaded.".format(info['Model']))
  101. model.trainable = False
  102. model.status = status
  103. return model
  104. def fix_input_shape(info, fixed_input_shape=None):
  105. if fixed_input_shape is not None:
  106. resize = {'ResizeByShort': {}}
  107. padding = {'Padding': {}}
  108. if info['_Attributes']['model_type'] == 'classifier':
  109. pass
  110. else:
  111. resize['ResizeByShort']['short_size'] = min(fixed_input_shape)
  112. resize['ResizeByShort']['max_size'] = max(fixed_input_shape)
  113. padding['Padding']['target_size'] = list(fixed_input_shape)
  114. info['Transforms'].append(resize)
  115. info['Transforms'].append(padding)