distillationloss.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  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 paddle
  15. import paddle.nn as nn
  16. from .celoss import CELoss
  17. from .dmlloss import DMLLoss
  18. from .distanceloss import DistanceLoss
  19. class DistillationCELoss(CELoss):
  20. """
  21. DistillationCELoss
  22. """
  23. def __init__(self,
  24. model_name_pairs=[],
  25. epsilon=None,
  26. key=None,
  27. name="loss_ce"):
  28. super().__init__(epsilon=epsilon)
  29. assert isinstance(model_name_pairs, list)
  30. self.key = key
  31. self.model_name_pairs = model_name_pairs
  32. self.name = name
  33. def forward(self, predicts, batch):
  34. loss_dict = dict()
  35. for idx, pair in enumerate(self.model_name_pairs):
  36. out1 = predicts[pair[0]]
  37. out2 = predicts[pair[1]]
  38. if self.key is not None:
  39. out1 = out1[self.key]
  40. out2 = out2[self.key]
  41. loss = super().forward(out1, out2)
  42. for key in loss:
  43. loss_dict["{}_{}_{}".format(key, pair[0], pair[1])] = loss[key]
  44. return loss_dict
  45. class DistillationGTCELoss(CELoss):
  46. """
  47. DistillationGTCELoss
  48. """
  49. def __init__(self,
  50. model_names=[],
  51. epsilon=None,
  52. key=None,
  53. name="loss_gt_ce"):
  54. super().__init__(epsilon=epsilon)
  55. assert isinstance(model_names, list)
  56. self.key = key
  57. self.model_names = model_names
  58. self.name = name
  59. def forward(self, predicts, batch):
  60. loss_dict = dict()
  61. for idx, name in enumerate(self.model_names):
  62. out = predicts[name]
  63. if self.key is not None:
  64. out = out[self.key]
  65. loss = super().forward(out, batch)
  66. for key in loss:
  67. loss_dict["{}_{}".format(key, name)] = loss[key]
  68. return loss_dict
  69. class DistillationDMLLoss(DMLLoss):
  70. """
  71. """
  72. def __init__(self,
  73. model_name_pairs=[],
  74. act=None,
  75. key=None,
  76. name="loss_dml"):
  77. super().__init__(act=act)
  78. assert isinstance(model_name_pairs, list)
  79. self.key = key
  80. self.model_name_pairs = model_name_pairs
  81. self.name = name
  82. def forward(self, predicts, batch):
  83. loss_dict = dict()
  84. for idx, pair in enumerate(self.model_name_pairs):
  85. out1 = predicts[pair[0]]
  86. out2 = predicts[pair[1]]
  87. if self.key is not None:
  88. out1 = out1[self.key]
  89. out2 = out2[self.key]
  90. loss = super().forward(out1, out2)
  91. if isinstance(loss, dict):
  92. for key in loss:
  93. loss_dict["{}_{}_{}_{}".format(key, pair[0], pair[1],
  94. idx)] = loss[key]
  95. else:
  96. loss_dict["{}_{}".format(self.name, idx)] = loss
  97. return loss_dict
  98. class DistillationDistanceLoss(DistanceLoss):
  99. """
  100. """
  101. def __init__(self,
  102. mode="l2",
  103. model_name_pairs=[],
  104. key=None,
  105. name="loss_",
  106. **kargs):
  107. super().__init__(mode=mode, **kargs)
  108. assert isinstance(model_name_pairs, list)
  109. self.key = key
  110. self.model_name_pairs = model_name_pairs
  111. self.name = name + "_l2"
  112. def forward(self, predicts, batch):
  113. loss_dict = dict()
  114. for idx, pair in enumerate(self.model_name_pairs):
  115. out1 = predicts[pair[0]]
  116. out2 = predicts[pair[1]]
  117. if self.key is not None:
  118. out1 = out1[self.key]
  119. out2 = out2[self.key]
  120. loss = super().forward(out1, out2)
  121. for key in loss:
  122. loss_dict["{}_{}_{}".format(self.name, key, idx)] = loss[key]
  123. return loss_dict