manager.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  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 inspect
  15. from collections.abc import Sequence
  16. class ComponentManager:
  17. """
  18. Implement a manager class to add the new component properly.
  19. The component can be added as either class or function type.
  20. Args:
  21. name (str): The name of component.
  22. Returns:
  23. A callable object of ComponentManager.
  24. Examples 1:
  25. from paddlex.paddleseg.cvlibs.manager import ComponentManager
  26. model_manager = ComponentManager()
  27. class AlexNet: ...
  28. class ResNet: ...
  29. model_manager.add_component(AlexNet)
  30. model_manager.add_component(ResNet)
  31. # Or pass a sequence alliteratively:
  32. model_manager.add_component([AlexNet, ResNet])
  33. print(model_manager.components_dict)
  34. # {'AlexNet': <class '__main__.AlexNet'>, 'ResNet': <class '__main__.ResNet'>}
  35. Examples 2:
  36. # Or an easier way, using it as a Python decorator, while just add it above the class declaration.
  37. from paddlex.paddleseg.cvlibs.manager import ComponentManager
  38. model_manager = ComponentManager()
  39. @model_manager.add_component
  40. class AlexNet: ...
  41. @model_manager.add_component
  42. class ResNet: ...
  43. print(model_manager.components_dict)
  44. # {'AlexNet': <class '__main__.AlexNet'>, 'ResNet': <class '__main__.ResNet'>}
  45. """
  46. def __init__(self, name=None):
  47. self._components_dict = dict()
  48. self._name = name
  49. def __len__(self):
  50. return len(self._components_dict)
  51. def __repr__(self):
  52. name_str = self._name if self._name else self.__class__.__name__
  53. return "{}:{}".format(name_str, list(self._components_dict.keys()))
  54. def __getitem__(self, item):
  55. if item not in self._components_dict.keys():
  56. raise KeyError("{} does not exist in availabel {}".format(item,
  57. self))
  58. return self._components_dict[item]
  59. @property
  60. def components_dict(self):
  61. return self._components_dict
  62. @property
  63. def name(self):
  64. return self._name
  65. def _add_single_component(self, component):
  66. """
  67. Add a single component into the corresponding manager.
  68. Args:
  69. component (function|class): A new component.
  70. Raises:
  71. TypeError: When `component` is neither class nor function.
  72. KeyError: When `component` was added already.
  73. """
  74. # Currently only support class or function type
  75. if not (inspect.isclass(component) or inspect.isfunction(component)):
  76. raise TypeError("Expect class/function type, but received {}".
  77. format(type(component)))
  78. # Obtain the internal name of the component
  79. component_name = component.__name__
  80. # Check whether the component was added already
  81. if component_name in self._components_dict.keys():
  82. raise KeyError("{} exists already!".format(component_name))
  83. else:
  84. # Take the internal name of the component as its key
  85. self._components_dict[component_name] = component
  86. def add_component(self, components):
  87. """
  88. Add component(s) into the corresponding manager.
  89. Args:
  90. components (function|class|list|tuple): Support four types of components.
  91. Returns:
  92. components (function|class|list|tuple): Same with input components.
  93. """
  94. # Check whether the type is a sequence
  95. if isinstance(components, Sequence):
  96. for component in components:
  97. self._add_single_component(component)
  98. else:
  99. component = components
  100. self._add_single_component(component)
  101. return components
  102. MODELS = ComponentManager("models")
  103. BACKBONES = ComponentManager("backbones")
  104. DATASETS = ComponentManager("datasets")
  105. TRANSFORMS = ComponentManager("transforms")
  106. LOSSES = ComponentManager("losses")