load_model.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  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. def load_model(model_dir):
  24. if not osp.exists(osp.join(model_dir, "model.yml")):
  25. raise Exception("There's not model.yml in {}".format(model_dir))
  26. with open(osp.join(model_dir, "model.yml")) as f:
  27. info = yaml.load(f.read(), Loader=yaml.Loader)
  28. status = info['status']
  29. if not hasattr(paddlex.cv.models, info['Model']):
  30. raise Exception("There's no attribute {} in paddlex.cv.models".format(
  31. info['Model']))
  32. if info['_Attributes']['model_type'] == 'classifier':
  33. model = paddlex.cv.models.BaseClassifier(**info['_init_params'])
  34. else:
  35. model = getattr(paddlex.cv.models,
  36. info['Model'])(**info['_init_params'])
  37. if status == "Normal" or \
  38. status == "Prune":
  39. startup_prog = fluid.Program()
  40. model.test_prog = fluid.Program()
  41. with fluid.program_guard(model.test_prog, startup_prog):
  42. with fluid.unique_name.guard():
  43. model.test_inputs, model.test_outputs = model.build_net(
  44. mode='test')
  45. model.test_prog = model.test_prog.clone(for_test=True)
  46. model.exe.run(startup_prog)
  47. if status == "Prune":
  48. from .slim.prune import update_program
  49. model.test_prog = update_program(model.test_prog, model_dir,
  50. model.places[0])
  51. import pickle
  52. with open(osp.join(model_dir, 'model.pdparams'), 'rb') as f:
  53. load_dict = pickle.load(f)
  54. fluid.io.set_program_state(model.test_prog, load_dict)
  55. elif status == "Infer" or \
  56. status == "Quant":
  57. [prog, input_names, outputs] = fluid.io.load_inference_model(
  58. model_dir, model.exe, params_filename='__params__')
  59. model.test_prog = prog
  60. test_outputs_info = info['_ModelInputsOutputs']['test_outputs']
  61. model.test_inputs = OrderedDict()
  62. model.test_outputs = OrderedDict()
  63. for name in input_names:
  64. model.test_inputs[name] = model.test_prog.global_block().var(name)
  65. for i, out in enumerate(outputs):
  66. var_desc = test_outputs_info[i]
  67. model.test_outputs[var_desc[0]] = out
  68. if 'Transforms' in info:
  69. transforms_mode = info.get('TransformsMode', 'RGB')
  70. if transforms_mode == 'RGB':
  71. to_rgb = True
  72. else:
  73. to_rgb = False
  74. model.test_transforms = build_transforms(model.model_type,
  75. info['Transforms'], to_rgb)
  76. model.eval_transforms = copy.deepcopy(model.test_transforms)
  77. if '_Attributes' in info:
  78. for k, v in info['_Attributes'].items():
  79. if k in model.__dict__:
  80. model.__dict__[k] = v
  81. logging.info("Model[{}] loaded.".format(info['Model']))
  82. return model
  83. def build_transforms(model_type, transforms_info, to_rgb=True):
  84. if model_type == "classifier":
  85. import paddlex.cv.transforms.cls_transforms as T
  86. elif model_type == "detector":
  87. import paddlex.cv.transforms.det_transforms as T
  88. elif model_type == "segmenter":
  89. import paddlex.cv.transforms.seg_transforms as T
  90. transforms = list()
  91. for op_info in transforms_info:
  92. op_name = list(op_info.keys())[0]
  93. op_attr = op_info[op_name]
  94. if not hasattr(T, op_name):
  95. raise Exception(
  96. "There's no operator named '{}' in transforms of {}".format(
  97. op_name, model_type))
  98. transforms.append(getattr(T, op_name)(**op_attr))
  99. eval_transforms = T.Compose(transforms)
  100. eval_transforms.to_rgb = to_rgb
  101. return eval_transforms