doc

nnet/loss.py

@@ -251,6 +251,11 @@ class ArcMarginProduct(torch.nn.Module):
         self.mm = math.sin(math.pi - self.m) * self.m
 
     def change_params(self, s=None, m=None):
+        """
+
+        :param s:
+        :param m:
+        """
         if s is None:
             s = self.s
         if m is None:
@@ -263,8 +268,15 @@ class ArcMarginProduct(torch.nn.Module):
         self.mm = math.sin(math.pi - self.m) * self.m
 
     def forward(self, input, target=None):
+        """
+
+        :param input:
+        :param target:
+        :return:
+        """
         # cos(theta)
-        cosine = torch.nn.functional.linear(torch.nn.functional.normalize(input), torch.nn.functional.normalize(self.weight))
+        cosine = torch.nn.functional.linear(torch.nn.functional.normalize(input),
+                                        torch.nn.functional.normalize(self.weight))
         if target == None:
             return cosine * self.s
         # cos(theta + m)