title: Focal loss implementation
body: <p>In the <a href="https://arxiv.org/pdf/1708.02002.pdf" rel="nofollow noreferrer">paper</a> introducing focal loss, they state that the loss function is formulated as such:</p>
  <p><a href="https://i.stack.imgur.com/s6Oo2.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/s6Oo2.png" alt="enter image description here" /></a></p>
  <p>Where</p>
  <p><a href="https://i.stack.imgur.com/5uMhW.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/5uMhW.png" alt="enter image description here" /></a></p>
  <p>I found an implementation of it on a Github page from another author who used it in their <a href="https://arxiv.org/pdf/2009.13935.pdf" rel="nofollow noreferrer">paper</a>. I tried the function out on a segmentation problem dataset I have and it seems to work quite well.</p>
  <p>Below is the implementation:</p>
  <pre><code>def binary_focal_loss(pred, truth, gamma=2., alpha=.25):
      eps = 1e-8
      pred = nn.Softmax(1)(pred)
      truth = F.one_hot(truth, num_classes = pred.shape[1]).permute(0,3,1,2).contiguous()
  
      pt_1 = torch.where(truth == 1, pred, torch.ones_like(pred))
      pt_0 = torch.where(truth == 0, pred, torch.zeros_like(pred))
  
      pt_1 = torch.clamp(pt_1, eps, 1. - eps)
      pt_0 = torch.clamp(pt_0, eps, 1. - eps)
  
      out1 = -torch.mean(alpha * torch.pow(1. - pt_1, gamma) * torch.log(pt_1)) 
      out0 = -torch.mean((1 - alpha) * torch.pow(pt_0, gamma) * torch.log(1. - pt_0))
  
      return out1 + out0
  </code></pre>
  <p>The part that I don't understand is the calculation of pt_0 and pt_1. I created a small example for myself to try and figure it out but it still confuses me a bit.</p>
  <pre><code># one hot encoded prediction tensor
  pred = torch.tensor([
                       [
                        [.2, .7, .8], # probability
                        [.3, .5, .7], # of
                        [.2, .6, .5]  # background class
                       ], 
                       [
                        [.8, .3, .2], # probability
                        [.7, .5, .3], # of
                        [.8, .4, .5]  # class 1
                       ]
                      ])
  
  # one-hot encoded ground truth labels
  truth = torch.tensor([
                        [1, 0, 0], 
                        [1, 1, 0], 
                        [1, 0, 0]
                       ])
  truth = F.one_hot(truth, num_classes = 2).permute(2,0,1).contiguous()
  
  print(truth)
  # gives me:
  # tensor([
  #         [
  #          [0, 1, 1],
  #          [0, 0, 1],
  #          [0, 1, 1]
  #         ],
  #         [
  #          [1, 0, 0],
  #          [1, 1, 0],
  #          [1, 0, 0]
  #         ]
  #       ])
  
  pt_0 = torch.where(truth == 0, pred, torch.zeros_like(pred))
  pt_1 = torch.where(truth == 1, pred, torch.ones_like(pred))
  
  print(pt_0)
  # gives me:
  # tensor([[
  #         [0.2000, 0.0000, 0.0000],
  #         [0.3000, 0.5000, 0.0000],
  #         [0.2000, 0.0000, 0.0000]
  #         ],
  #        [
  #         [0.0000, 0.3000, 0.2000],
  #         [0.0000, 0.0000, 0.3000],
  #         [0.0000, 0.4000, 0.5000]
  #        ]
  #      ])
  
  print(pt_1)
  # gives me:
  # tensor([[
  #          [1.0000, 0.7000, 0.8000],
  #          [1.0000, 1.0000, 0.7000],
  #          [1.0000, 0.6000, 0.5000]
  #         ],
  #         [
  #          [0.8000, 1.0000, 1.0000],
  #          [0.7000, 0.5000, 1.0000],
  #          [0.8000, 1.0000, 1.0000]
  #         ]
  #       ])
  </code></pre>
  <p>What I don't understand is why in pt_0 we are placing zeros where the torch.where statement is false, and in pt_1 we place ones. From how I understood the paper, I would have thought that instead of placing zeros or ones, you would place 1-p.</p>
  <p>Can anyone help explain this to me?</p>