fix stability in bce with pos_weight formula (#13863)

Ailing Zhang · facebook-github-bot · commit a17c0118a52d · 2018-11-12T22:04:24.000-08:00
Summary: Fixes #13773 Pull Request resolved: #13863 Differential Revision: D13031803 Pulled By: ailzhang fbshipit-source-id: 6c9e044f0450eebf4555bbc02c125713d9378e2f
diff --git a/aten/src/ATen/native/Loss.cpp b/aten/src/ATen/native/Loss.cpp
@@ -93,7 +93,7 @@ Tensor binary_cross_entropy_with_logits(const Tensor& input, const Tensor& targe
     if (pos_weight.defined()) {
         // pos_weight need to be broadcasted, thus mul(target) is not inplace.
         auto log_weight = (pos_weight - 1).mul(target).add_(1);
-        loss = (1 - target).mul_(input).add_(log_weight.mul_((-max_val).exp_().mul_(1 + (-input).exp_()).log_().add_(max_val)));
+        loss = (1 - target).mul_(input).add_(log_weight.mul_(((-max_val).exp_().add_((-input - max_val).exp_())).log_().add_(max_val)));
     } else {
         loss = (1 - target).mul_(input).add_(max_val).add_((-max_val).exp_().add_((-input -max_val).exp_()).log_());
     }
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -4939,6 +4939,17 @@ def test_bce_with_logits_with_pos_weight_has_correct_grad_at_zero(self):
         grad = output.grad
         self.assertEqual(grad, expected_grad)
 
+    def test_bce_with_logits_stability(self):
+        output = torch.tensor([0., -120.])
+        target = torch.tensor([0., 1.])
+        pos_weight = torch.tensor([1., 1.])
+
+        out1 = nn.BCEWithLogitsLoss()(output, target)
+        self.assertTrue(torch.isfinite(out1).all().item())
+
+        out2 = nn.BCEWithLogitsLoss(pos_weight=pos_weight)(output, target)
+        self.assertTrue(torch.isfinite(out2).all().item())
+
     def test_bce_loss_broadcasts_weights(self):
         sigmoid = nn.Sigmoid()
         target = torch.rand(16, 4)

Original file line number	Diff line number	Diff line change
`@@ -93,7 +93,7 @@ Tensor binary_cross_entropy_with_logits(const Tensor& input, const Tensor& targe`
`93`	`93`	`if (pos_weight.defined()) {`
`94`	`94`	`// pos_weight need to be broadcasted, thus mul(target) is not inplace.`
`95`	`95`	`auto log_weight = (pos_weight - 1).mul(target).add_(1);`
`96`		`- loss = (1 - target).mul_(input).add_(log_weight.mul_((-max_val).exp_().mul_(1 + (-input).exp_()).log_().add_(max_val)));`
	`96`	`+ loss = (1 - target).mul_(input).add_(log_weight.mul_(((-max_val).exp_().add_((-input - max_val).exp_())).log_().add_(max_val)));`
`97`	`97`	`} else {`
`98`	`98`	`loss = (1 - target).mul_(input).add_(max_val).add_((-max_val).exp_().add_((-input -max_val).exp_()).log_());`
`99`	`99`	`}`