refactoring

Author: AdeelH

.gitignore

@@ -1,2 +1,4 @@
 
 *.pyc
+
+.vscode/

focal_loss.py

@@ -1,3 +1,5 @@
+from typing import Optional, Sequence
+
 import torch
 from torch import Tensor
 from torch import nn
@@ -18,17 +20,20 @@ class FocalLoss(nn.Module):
     """
 
     def __init__(self,
-                 alpha: Tensor = None,
+                 alpha: Optional[Tensor] = None,
                  gamma: float = 0.,
                  reduction: str = 'mean',
                  ignore_index: int = -100):
         """Constructor.
+
         Args:
-            alpha (Tensor): Weights for each class.
-            gamma (float): A constant, as described in the paper.
+            alpha (Tensor, optional): Weights for each class. Defaults to None.
+            gamma (float, optional): A constant, as described in the paper.
+                Defaults to 0.
             reduction (str, optional): 'mean', 'sum' or 'none'.
                 Defaults to 'mean'.
             ignore_index (int, optional): class label to ignore.
+                Defaults to -100.
         """
         if reduction not in ('mean', 'sum', 'none'):
             raise ValueError(
@@ -65,6 +70,7 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor:
         x = x[unignored_mask]
 
         # compute weighted cross entropy term: -alpha * log(pt)
+        # (alpha is already part of self.nll_loss)
         log_p = F.log_softmax(x, dim=-1)
         ce = self.nll_loss(log_p, y)
 
@@ -87,15 +93,38 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor:
         return loss
 
 
-def focal_loss(alpha=None, gamma=0., reduction='mean', ignore_index=-100,
-				device='cpu', dtype=torch.float32):
-	if not ((alpha is None) or isinstance(alpha, torch.Tensor)):
-		alpha = torch.tensor(alpha, device=device, dtype=dtype)
-
-	fl = FocalLoss(
-		alpha=alpha,
-		gamma=gamma,
-		reduction=reduction,
-		ignore_index=ignore_index
-	)
-	return fl
+def focal_loss(alpha: Optional[Sequence] = None,
+               gamma: float = 0.,
+               reduction: str = 'mean',
+               ignore_index: int = -100,
+               device='cpu',
+               dtype=torch.float32) -> FocalLoss:
+    """Factory function for FocalLoss.
+
+    Args:
+        alpha (Sequence, optional): Weights for each class. Will be converted
+            to a Tensor if not None. Defaults to None.
+        gamma (float, optional): A constant, as described in the paper.
+            Defaults to 0.
+        reduction (str, optional): 'mean', 'sum' or 'none'.
+            Defaults to 'mean'.
+        ignore_index (int, optional): class label to ignore.
+            Defaults to -100.
+        device (str, optional): Device to move alpha to. Defaults to 'cpu'.
+        dtype (torch.dtype, optional): dtype to cast alpha to.
+            Defaults to torch.float32.
+
+    Returns:
+        A FocalLoss object
+    """
+    if alpha is not None:
+        if not isinstance(alpha, Tensor):
+            alpha = torch.tensor(alpha)
+        alpha = alpha.to(device=device, dtype=dtype)
+
+    fl = FocalLoss(
+        alpha=alpha,
+        gamma=gamma,
+        reduction=reduction,
+        ignore_index=ignore_index)
+    return fl