theseus_layer.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. from abc import ABC
  2. from paddle import nn
  3. import re
  4. class Identity(nn.Layer):
  5. def __init__(self):
  6. super(Identity, self).__init__()
  7. def forward(self, inputs):
  8. return inputs
  9. class TheseusLayer(nn.Layer):
  10. def __init__(self, *args, **kwargs):
  11. super(TheseusLayer, self).__init__()
  12. self.res_dict = {}
  13. self.res_name = self.full_name()
  14. # stop doesn't work when stop layer has a parallel branch.
  15. def stop_after(self, stop_layer_name: str):
  16. after_stop = False
  17. for layer_i in self._sub_layers:
  18. if after_stop:
  19. self._sub_layers[layer_i] = Identity()
  20. continue
  21. layer_name = self._sub_layers[layer_i].full_name()
  22. if layer_name == stop_layer_name:
  23. after_stop = True
  24. continue
  25. if isinstance(self._sub_layers[layer_i], TheseusLayer):
  26. after_stop = self._sub_layers[layer_i].stop_after(
  27. stop_layer_name)
  28. return after_stop
  29. def update_res(self, return_patterns):
  30. for return_pattern in return_patterns:
  31. pattern_list = return_pattern.split(".")
  32. if not pattern_list:
  33. continue
  34. sub_layer_parent = self
  35. while len(pattern_list) > 1:
  36. if '[' in pattern_list[0]:
  37. sub_layer_name = pattern_list[0].split('[')[0]
  38. sub_layer_index = pattern_list[0].split('[')[1].split(']')[
  39. 0]
  40. sub_layer_parent = getattr(sub_layer_parent,
  41. sub_layer_name)[sub_layer_index]
  42. else:
  43. sub_layer_parent = getattr(sub_layer_parent,
  44. pattern_list[0], None)
  45. if sub_layer_parent is None:
  46. break
  47. if isinstance(sub_layer_parent, WrapLayer):
  48. sub_layer_parent = sub_layer_parent.sub_layer
  49. pattern_list = pattern_list[1:]
  50. if sub_layer_parent is None:
  51. continue
  52. if '[' in pattern_list[0]:
  53. sub_layer_name = pattern_list[0].split('[')[0]
  54. sub_layer_index = pattern_list[0].split('[')[1].split(']')[0]
  55. sub_layer = getattr(sub_layer_parent,
  56. sub_layer_name)[sub_layer_index]
  57. if not isinstance(sub_layer, TheseusLayer):
  58. sub_layer = wrap_theseus(sub_layer)
  59. getattr(sub_layer_parent,
  60. sub_layer_name)[sub_layer_index] = sub_layer
  61. else:
  62. sub_layer = getattr(sub_layer_parent, pattern_list[0])
  63. if not isinstance(sub_layer, TheseusLayer):
  64. sub_layer = wrap_theseus(sub_layer)
  65. setattr(sub_layer_parent, pattern_list[0], sub_layer)
  66. sub_layer.res_dict = self.res_dict
  67. sub_layer.res_name = return_pattern
  68. sub_layer.register_forward_post_hook(sub_layer._save_sub_res_hook)
  69. def _save_sub_res_hook(self, layer, input, output):
  70. self.res_dict[self.res_name] = output
  71. def _return_dict_hook(self, layer, input, output):
  72. res_dict = {"output": output}
  73. for res_key in list(self.res_dict):
  74. res_dict[res_key] = self.res_dict.pop(res_key)
  75. return res_dict
  76. def replace_sub(self, layer_name_pattern, replace_function,
  77. recursive=True):
  78. for layer_i in self._sub_layers:
  79. layer_name = self._sub_layers[layer_i].full_name()
  80. if re.match(layer_name_pattern, layer_name):
  81. self._sub_layers[layer_i] = replace_function(self._sub_layers[
  82. layer_i])
  83. if recursive:
  84. if isinstance(self._sub_layers[layer_i], TheseusLayer):
  85. self._sub_layers[layer_i].replace_sub(
  86. layer_name_pattern, replace_function, recursive)
  87. elif isinstance(self._sub_layers[layer_i],
  88. (nn.Sequential, nn.LayerList)):
  89. for layer_j in self._sub_layers[layer_i]._sub_layers:
  90. self._sub_layers[layer_i]._sub_layers[
  91. layer_j].replace_sub(layer_name_pattern,
  92. replace_function, recursive)
  93. '''
  94. example of replace function:
  95. def replace_conv(origin_conv: nn.Conv2D):
  96. new_conv = nn.Conv2D(
  97. in_channels=origin_conv._in_channels,
  98. out_channels=origin_conv._out_channels,
  99. kernel_size=origin_conv._kernel_size,
  100. stride=2
  101. )
  102. return new_conv
  103. '''
  104. class WrapLayer(TheseusLayer):
  105. def __init__(self, sub_layer):
  106. super(WrapLayer, self).__init__()
  107. self.sub_layer = sub_layer
  108. def forward(self, *inputs, **kwargs):
  109. return self.sub_layer(*inputs, **kwargs)
  110. def wrap_theseus(sub_layer):
  111. wrapped_layer = WrapLayer(sub_layer)
  112. return wrapped_layer