diff --git a/ppcls/arch/head/arcmargin.py b/ppcls/arch/head/arcmargin.py new file mode 100644 index 0000000000000000000000000000000000000000..82da7f093fde2949d1d5c15195318737550d4cac --- /dev/null +++ b/ppcls/arch/head/arcmargin.py @@ -0,0 +1,65 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import paddle +import paddle.nn as nn +import math + +class ArcMargin(nn.Layer): + def __init__(self, embedding_size, + class_num, + margin=0.5, + scale=80.0, + easy_margin=False): + super(ArcMargin, self).__init__() + self.embedding_size = embedding_size + self.class_num = class_num + self.margin = margin + self.scale = scale + self.easy_margin = easy_margin + + weight_attr = paddle.ParamAttr(initializer = paddle.nn.initializer.XavierNormal()) + self.fc = nn.Linear(self.embedding_size, self.class_dim, weight_attr=weight_attr, bias_attr=False) + + def forward(self, input, label): + input_norm = paddle.sqrt(paddle.sum(paddle.square(input), axis=1, keepdim=True)) + input = paddle.divide(input, input_norm) + + weight = self.fc0.weight + weight_norm = paddle.sqrt(paddle.sum(paddle.square(weight), axis=0, keepdim=True)) + weight = paddle.divide(weight, weight_norm) + + cos = paddle.matmul(input, weight) + sin = paddle.sqrt(1.0 - paddle.square(cos) + 1e-6) + cos_m = math.cos(self.margin) + sin_m = math.sin(self.margin) + phi = cos * cos_m - sin * sin_m + + th = math.cos(self.margin) * (-1) + mm = math.sin(self.margin) * self.margin + if self.easy_margin: + phi = self._paddle_where_more_than(cos, 0, phi, cos) + else: + phi = self._paddle_where_more_than(cos, th, phi, cos - mm) + + one_hot = paddle.nn.functional.one_hot(label, self.class_dim) + one_hot = paddle.squeeze(one_hot, axis=[1]) + output = paddle.multiply(one_hot, phi) + paddle.multiply((1.0 - one_hot), cos) + output = output * self.scale + return output + + def _paddle_where_more_than(self, target, limit, x, y): + mask = paddle.cast( x = (target > limit), dtype='float32') + output = paddle.multiply(mask, x) + paddle.multiply((1.0 - mask), y) + return output