__init__.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. #copyright (c) 2021 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 copy
  15. import importlib
  16. import paddle.nn as nn
  17. from paddle.jit import to_static
  18. from paddle.static import InputSpec
  19. from . import backbone, gears
  20. from .backbone import *
  21. from .gears import build_gear
  22. from .utils import *
  23. from paddlex.ppcls.arch.backbone.base.theseus_layer import TheseusLayer
  24. from paddlex.ppcls.utils import logger
  25. from paddlex.ppcls.utils.save_load import load_dygraph_pretrain
  26. __all__ = ["build_model", "RecModel", "DistillationModel"]
  27. def build_model(config):
  28. config = copy.deepcopy(config)
  29. model_type = config.pop("name")
  30. mod = importlib.import_module(__name__)
  31. arch = getattr(mod, model_type)(**config)
  32. return arch
  33. def apply_to_static(config, model):
  34. support_to_static = config['Global'].get('to_static', False)
  35. if support_to_static:
  36. specs = None
  37. if 'image_shape' in config['Global']:
  38. specs = [InputSpec([None] + config['Global']['image_shape'])]
  39. model = to_static(model, input_spec=specs)
  40. logger.info("Successfully to apply @to_static with specs: {}".format(
  41. specs))
  42. return model
  43. class RecModel(nn.Layer):
  44. def __init__(self, **config):
  45. super().__init__()
  46. backbone_config = config["Backbone"]
  47. backbone_name = backbone_config.pop("name")
  48. self.backbone = eval(backbone_name)(**backbone_config)
  49. if "BackboneStopLayer" in config:
  50. backbone_stop_layer = config["BackboneStopLayer"]["name"]
  51. self.backbone.stop_after(backbone_stop_layer)
  52. if "Neck" in config:
  53. self.neck = build_gear(config["Neck"])
  54. else:
  55. self.neck = None
  56. if "Head" in config:
  57. self.head = build_gear(config["Head"])
  58. else:
  59. self.head = None
  60. def forward(self, x, label=None):
  61. x = self.backbone(x)
  62. if self.neck is not None:
  63. x = self.neck(x)
  64. if self.head is not None:
  65. y = self.head(x, label)
  66. else:
  67. y = None
  68. return {"features": x, "logits": y}
  69. class DistillationModel(nn.Layer):
  70. def __init__(self,
  71. models=None,
  72. pretrained_list=None,
  73. freeze_params_list=None,
  74. **kargs):
  75. super().__init__()
  76. assert isinstance(models, list)
  77. self.model_list = []
  78. self.model_name_list = []
  79. if pretrained_list is not None:
  80. assert len(pretrained_list) == len(models)
  81. if freeze_params_list is None:
  82. freeze_params_list = [False] * len(models)
  83. assert len(freeze_params_list) == len(models)
  84. for idx, model_config in enumerate(models):
  85. assert len(model_config) == 1
  86. key = list(model_config.keys())[0]
  87. model_config = model_config[key]
  88. model_name = model_config.pop("name")
  89. model = eval(model_name)(**model_config)
  90. if freeze_params_list[idx]:
  91. for param in model.parameters():
  92. param.trainable = False
  93. self.model_list.append(self.add_sublayer(key, model))
  94. self.model_name_list.append(key)
  95. if pretrained_list is not None:
  96. for idx, pretrained in enumerate(pretrained_list):
  97. if pretrained is not None:
  98. load_dygraph_pretrain(
  99. self.model_name_list[idx], path=pretrained)
  100. def forward(self, x, label=None):
  101. result_dict = dict()
  102. for idx, model_name in enumerate(self.model_name_list):
  103. if label is None:
  104. result_dict[model_name] = self.model_list[idx](x)
  105. else:
  106. result_dict[model_name] = self.model_list[idx](x, label)
  107. return result_dict