tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 12:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

improve Tensor.cross_entropy (#10985)

Browse Source 
separate when Y is prob vs indices and check shapes for indices. also fix higher dim cases
This commit is contained in:
chenyu
2025-06-26 01:39:48 -04:00
committed by GitHub
parent 8751d47985
commit 0612acfc70
2 changed files with 21 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_ops.py
+15 -5
View File
@@ -2918,11 +2918,18 @@ class TestOps(unittest.TestCase):
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,y.clip(0,1),
pos_weight=torch.tensor(pos_weight)),
lambda x,y: x.binary_crossentropy_logits(y.clip(0,1),pos_weight=Tensor(pos_weight)))
def test_cross_entropy(self):
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y),
lambda x,y: x.cross_entropy(y))
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, torch.argmax(y, dim=1)),
lambda x,y: x.cross_entropy(y.argmax(axis=1)), forward_only=True)
def test_cross_entropy_class_probabilities(self):
helper_test_op([(32,), (32,)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
helper_test_op([(32,4,4,4), (32,4,4,4)], lambda x,y: torch.nn.functional.cross_entropy(x, y), lambda x,y: x.cross_entropy(y))
def test_cross_entropy_class_indices(self):
classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
helper_test_op([(32,10)], lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes)),
lambda x: x.cross_entropy(Tensor(classes)))
self.helper_test_exception([(32,10), (32,1)], lambda x,y: torch.nn.functional.cross_entropy(x, y),
lambda x,y: x.cross_entropy(y), expected=(AssertionError, RuntimeError))
def test_cross_entropy_reductions(self):
for r in ("mean", "sum", "none"):
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y, reduction=r),
@@ -2934,6 +2941,9 @@ class TestOps(unittest.TestCase):
for ls in (0., 0.3, 0.7, 1.):
helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.cross_entropy(x, y, label_smoothing=ls),
lambda x,y: x.cross_entropy(y, label_smoothing=ls))
classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
helper_test_op([(32,10)], lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes), label_smoothing=ls),
lambda x: x.cross_entropy(Tensor(classes), label_smoothing=ls))
def test_nll_loss(self):
helper_test_op([(32,10), (32)],
tinygrad/tensor.py
+6 -4
View File
@@ -3948,10 +3948,12 @@ class Tensor(MathTrait):
```
"""
assert 0.0 <= label_smoothing <= 1.0, "label_smoothing must be in [0.0, 1.0]"
Y = Y.one_hot(num_classes=cast(int, self.shape[1])) if Y.ndim < 2 else Y
Y = (1 - label_smoothing)*Y + label_smoothing / cast(int, Y.shape[1])
ret = -self.log_softmax(axis=1).mul(Y).sum(axis=1)
return ret._do_reduction(reduction)
classes_dim = 0 if self.ndim == 1 else 1
if self.shape != Y.shape:
if self.max(classes_dim).shape != Y.shape: raise RuntimeError(f"shape mismatch: {self.shape=}, {Y.shape=}")
Y = Y.unsqueeze(classes_dim)._one_hot_along_dim(num_classes=self.shape[classes_dim], dim=classes_dim)
Y = (1 - label_smoothing)*Y + label_smoothing / int(Y.shape[classes_dim])
return -self.log_softmax(classes_dim).mul(Y).sum(classes_dim)._do_reduction(reduction)
def nll_loss(self, Y:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Tensor:
"""