Repository: pyg-team/pytorch_geometric
Pull request #8340 (2023-11-07 16:12:03)

Added support for weighted `sparse_cross_entropy`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8340?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8340](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8340?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (27f0af4) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/a2b77e67e774aeaf7b638619d702b205f61ba7d9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (a2b77e6) will **decrease** coverage by `0.68%`.
> Report is 1 commits behind head on master.
> The diff coverage is `100.00%`.

> :exclamation: Current head 27f0af4 differs from pull request most recent head 019b0e0. Consider uploading reports for the commit 019b0e0 to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #8340      +/-   ##
==========================================
- Coverage   87.77%   87.10%   -0.68%     
==========================================
  Files         473      473              
  Lines       28751    28754       +3     
==========================================
- Hits        25237    25046     -191     
- Misses       3514     3708     +194     
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8340?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/cross\_entropy.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8340?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2Nyb3NzX2VudHJvcHkucHk=) | `100.00% <100.00%> (ø)` | |

... and [35 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8340/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: Codecov offers a browser extension for seamless coverage viewing on GitHub. Try it in [Chrome](https://chrome.google.com/webstore/detail/codecov/gedikamndpbemklijjkncpnolildpbgo) or [Firefox](https://addons.mozilla.org/en-US/firefox/addon/codecov/) today!


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b86a76145..881e376fe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for weighted `sparse_cross_entropy` ([#8340](https://github.com/pyg-team/pytorch_geometric/pull/8340))
 - Added a multi GPU training benchmarks for XPU device ([#8288](https://github.com/pyg-team/pytorch_geometric/pull/8288))
 - Support MRR computation in `KGEModel.test()` ([#8298](https://github.com/pyg-team/pytorch_geometric/pull/8298))
 - Added an example for model parallelism (`examples/multi_gpu/model_parallel.py`) ([#8309](https://github.com/pyg-team/pytorch_geometric/pull/8309))
diff --git a/torch_geometric/utils/cross_entropy.py b/torch_geometric/utils/cross_entropy.py
index 39b9e45e2..fc8180e20 100644
--- a/torch_geometric/utils/cross_entropy.py
+++ b/torch_geometric/utils/cross_entropy.py
@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -10,26 +10,38 @@ class SparseCrossEntropy(torch.autograd.Function):
     # We implement our own custom autograd function for this to avoid the
     # double gradient computation to `inputs`.
     @staticmethod
-    def forward(ctx, inputs: Tensor, edge_label_index: Tensor) -> Tensor:
+    def forward(
+        ctx,
+        inputs: Tensor,
+        edge_label_index: Tensor,
+        edge_label_weight: Optional[Tensor],
+    ) -> Tensor:
         assert inputs.dim() == 2
 
         logsumexp = inputs.logsumexp(dim=-1)
-        ctx.save_for_backward(inputs, edge_label_index, logsumexp)
+        ctx.save_for_backward(inputs, edge_label_index, edge_label_weight,
+                              logsumexp)
 
         out = inputs[edge_label_index[0], edge_label_index[1]]
         out.neg_().add_(logsumexp[edge_label_index[0]])
+        if edge_label_weight is not None:
+            out *= edge_label_weight
 
         return out.sum() / inputs.size(0)
 
     @staticmethod
     @torch.autograd.function.once_differentiable
-    def backward(ctx, grad_out: Tensor) -> Tuple[Tensor, None]:
-        inputs, edge_label_index, logsumexp = ctx.saved_tensors
+    def backward(ctx, grad_out: Tensor) -> Tuple[Tensor, None, None]:
+        inputs, edge_label_index, edge_label_weight, logsumexp = (
+            ctx.saved_tensors)
 
         grad_out = grad_out / inputs.size(0)
+        grad_out = grad_out.expand(edge_label_index.size(1))
 
-        grad_logsumexp = scatter(grad_out.expand(edge_label_index.size(1)),
-                                 edge_label_index[0], dim=0,
+        if edge_label_weight is not None:
+            grad_out = grad_out * edge_label_weight
+
+        grad_logsumexp = scatter(grad_out, edge_label_index[0], dim=0,
                                  dim_size=inputs.size(0), reduce='sum')
 
         # Gradient computation of `logsumexp`: `grad * (self - result).exp()`
@@ -39,10 +51,14 @@ class SparseCrossEntropy(torch.autograd.Function):
 
         grad_input[edge_label_index[0], edge_label_index[1]] -= grad_out
 
-        return grad_input, None
+        return grad_input, None, None
 
 
-def sparse_cross_entropy(inputs: Tensor, edge_label_index: Tensor) -> Tensor:
+def sparse_cross_entropy(
+    inputs: Tensor,
+    edge_label_index: Tensor,
+    edge_label_weight: Optional[Tensor] = None,
+) -> Tensor:
     r"""A sparse-label variant of :func:`torch.nn.functional.cross_entropy`.
     In particular, the binary target matrix is solely given by sparse indices
     :obj:`edge_label_index`.
@@ -50,16 +66,27 @@ def sparse_cross_entropy(inputs: Tensor, edge_label_index: Tensor) -> Tensor:
     Args:
         inputs (torch.Tensor): The predicted unnormalized logits of shape
             :obj:`[batch_size, num_classes]`.
-        edge_label_index (torch.Tensor): The sparse ground-truth indices of
+        edge_label_index (torch.Tensor): The sparse ground-truth indices with
             shape :obj:`[2, num_labels]`.
+        edge_label_weight (torch.Tensor, optional): The weight of ground-truth
+            indices with shape :obj:`[num_labels]`. (default: :obj:`None`)
 
     :rtype: :class:`torch.Tensor`
 
     Example:
         >>> inputs = torch.randn(2, 3)
-        >>> edge_label_index = torch.tensor([[0, 0, 1],
-        ...                                  [0, 1, 2]])
-        >>> sparse_cross_entropy(inputs, edge_label_index)
+        >>> edge_label_index = torch.tensor([
+        ...     [0, 0, 1],
+        ...     [0, 1, 2],
+        ... ])
+        >>> loss = sparse_cross_entropy(inputs, edge_label_index)
         tensor(1.2919)
     """
-    return SparseCrossEntropy.apply(inputs, edge_label_index)
+    if edge_label_weight is not None:
+        assert not edge_label_weight.requires_grad
+
+    return SparseCrossEntropy.apply(
+        inputs,
+        edge_label_index,
+        edge_label_weight,
+    )


Test patch:
diff --git a/test/utils/test_cross_entropy.py b/test/utils/test_cross_entropy.py
index 4782cb4c7..fca33519b 100644
--- a/test/utils/test_cross_entropy.py
+++ b/test/utils/test_cross_entropy.py
@@ -1,37 +1,50 @@
+import pytest
 import torch
 import torch.nn.functional as F
 
 from torch_geometric.utils.cross_entropy import sparse_cross_entropy
 
 
-def test_sparse_cross_entropy_multiclass():
+@pytest.mark.parametrize('with_edge_label_weight', [False, True])
+def test_sparse_cross_entropy_multiclass(with_edge_label_weight):
     x = torch.randn(5, 5, requires_grad=True)
     y = torch.eye(5)
+
     edge_label_index = y.nonzero().t()
+    edge_label_weight = None
+    if with_edge_label_weight:
+        edge_label_weight = torch.rand(edge_label_index.size(1))
+        y[y == 1.0] = edge_label_weight
 
     expected = F.cross_entropy(x, y)
     expected.backward()
     expected_grad = x.grad
 
     x.grad = None
-    out = sparse_cross_entropy(x, edge_label_index)
+    out = sparse_cross_entropy(x, edge_label_index, edge_label_weight)
     out.backward()
 
     assert torch.allclose(expected, out)
     assert torch.allclose(expected_grad, x.grad)
 
 
-def test_sparse_cross_entropy_multilabel():
+@pytest.mark.parametrize('with_edge_label_weight', [False, True])
+def test_sparse_cross_entropy_multilabel(with_edge_label_weight):
     x = torch.randn(4, 4, requires_grad=True)
     y = torch.randint_like(x, 0, 2)
+
     edge_label_index = y.nonzero().t()
+    edge_label_weight = None
+    if with_edge_label_weight:
+        edge_label_weight = torch.rand(edge_label_index.size(1))
+        y[y == 1.0] = edge_label_weight
 
     expected = F.cross_entropy(x, y)
     expected.backward()
     expected_grad = x.grad
 
     x.grad = None
-    out = sparse_cross_entropy(x, edge_label_index)
+    out = sparse_cross_entropy(x, edge_label_index, edge_label_weight)
     out.backward()
 
     assert torch.allclose(expected, out)