arcmargin.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. # Copyright (c) 2021 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 paddle
  15. import paddle.nn as nn
  16. import math
  17. class ArcMargin(nn.Layer):
  18. def __init__(self,
  19. embedding_size,
  20. class_num,
  21. margin=0.5,
  22. scale=80.0,
  23. easy_margin=False):
  24. super().__init__()
  25. self.embedding_size = embedding_size
  26. self.class_num = class_num
  27. self.margin = margin
  28. self.scale = scale
  29. self.easy_margin = easy_margin
  30. self.weight = self.create_parameter(
  31. shape=[self.embedding_size, self.class_num],
  32. is_bias=False,
  33. default_initializer=paddle.nn.initializer.XavierNormal())
  34. def forward(self, input, label=None):
  35. input_norm = paddle.sqrt(
  36. paddle.sum(paddle.square(input), axis=1, keepdim=True))
  37. input = paddle.divide(input, input_norm)
  38. weight_norm = paddle.sqrt(
  39. paddle.sum(paddle.square(self.weight), axis=0, keepdim=True))
  40. weight = paddle.divide(self.weight, weight_norm)
  41. cos = paddle.matmul(input, weight)
  42. if not self.training or label is None:
  43. return cos
  44. sin = paddle.sqrt(1.0 - paddle.square(cos) + 1e-6)
  45. cos_m = math.cos(self.margin)
  46. sin_m = math.sin(self.margin)
  47. phi = cos * cos_m - sin * sin_m
  48. th = math.cos(self.margin) * (-1)
  49. mm = math.sin(self.margin) * self.margin
  50. if self.easy_margin:
  51. phi = self._paddle_where_more_than(cos, 0, phi, cos)
  52. else:
  53. phi = self._paddle_where_more_than(cos, th, phi, cos - mm)
  54. one_hot = paddle.nn.functional.one_hot(label, self.class_num)
  55. one_hot = paddle.squeeze(one_hot, axis=[1])
  56. output = paddle.multiply(one_hot, phi) + paddle.multiply(
  57. (1.0 - one_hot), cos)
  58. output = output * self.scale
  59. return output
  60. def _paddle_where_more_than(self, target, limit, x, y):
  61. mask = paddle.cast(x=(target > limit), dtype='float32')
  62. output = paddle.multiply(mask, x) + paddle.multiply((1.0 - mask), y)
  63. return output