dmlloss.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  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. import paddle.nn.functional as F
  17. class DMLLoss(nn.Layer):
  18. """
  19. DMLLoss
  20. """
  21. def __init__(self, act="softmax"):
  22. super().__init__()
  23. if act is not None:
  24. assert act in ["softmax", "sigmoid"]
  25. if act == "softmax":
  26. self.act = nn.Softmax(axis=-1)
  27. elif act == "sigmoid":
  28. self.act = nn.Sigmoid()
  29. else:
  30. self.act = None
  31. def forward(self, out1, out2):
  32. if self.act is not None:
  33. out1 = self.act(out1)
  34. out2 = self.act(out2)
  35. log_out1 = paddle.log(out1)
  36. log_out2 = paddle.log(out2)
  37. loss = (F.kl_div(
  38. log_out1, out2, reduction='batchmean') + F.kl_div(
  39. log_out2, out1, reduction='batchmean')) / 2.0
  40. return {"DMLLoss": loss}