diff --git a/test/test_nn.py b/test/test_nn.py index 72fc0d77df..f85aa43f59 100755 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -13,29 +13,6 @@ from test.helpers import not_support_multi_device @unittest.skipIf(CI and Device.DEFAULT in {"CUDA", "NV"}, "slow") class TestNN(unittest.TestCase): - def test_sparse_categorical_crossentropy(self): - # create in tinygrad - input_tensor = Tensor.randn(6, 5) # not square to test that mean scaling uses the correct dimension - target = Tensor([0, 0, 0, 1, 2, 3]) # torch doesn't support target=-1 - torch_input = torch.tensor(input_tensor.numpy()) - torch_target = torch.tensor(target.numpy(), dtype=torch.long) - - for smoothing in [0.0, 0.1, 0.5, 1.0]: - for ignore_index in [-1, 0, 2]: - for reduction in ["none", "sum", "mean"]: - loss = input_tensor.sparse_categorical_crossentropy(target, label_smoothing=smoothing, ignore_index=ignore_index, reduction=reduction) - torch_loss = torch.nn.CrossEntropyLoss(reduction=reduction, label_smoothing=smoothing, ignore_index=ignore_index)(torch_input, torch_target) - np.testing.assert_allclose(loss.numpy(), torch_loss.detach().numpy(), atol=1e-5, rtol=1e-6) - - # also test with a batch dimension (of size 1) - loss = input_tensor.unsqueeze(0).sparse_categorical_crossentropy( - target.unsqueeze(0), label_smoothing=smoothing, ignore_index=ignore_index, reduction=reduction - ) - torch_loss = torch.nn.CrossEntropyLoss(reduction=reduction, label_smoothing=smoothing, ignore_index=ignore_index)( - torch_input.unsqueeze(0).permute(0,2,1), torch_target.unsqueeze(0) - ) - np.testing.assert_allclose(loss.numpy(), torch_loss.detach().numpy(), atol=1e-5, rtol=1e-6) - def test_batchnorm2d(self, training=False, threed=False, track_running_stats=True): with Tensor.train(training): szs = [4, 8, 16, 32] diff --git a/test/test_ops.py b/test/test_ops.py index cb94a4294f..6b6ef9563e 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -2950,6 +2950,39 @@ class TestOps(unittest.TestCase): 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_sparse_categorical_crossentropy(self): + classes = np.random.randint(0, 10, (12,), dtype=np.int32).tolist() + helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss()(x, torch.tensor(classes)), + lambda x: x.sparse_categorical_crossentropy(Tensor(classes))) + + # combine args + helper_test_op([(12,10)], + lambda x: torch.nn.CrossEntropyLoss(reduction="mean", ignore_index=classes[0], label_smoothing=0.3)(x, torch.tensor(classes)), + lambda x: x.sparse_categorical_crossentropy(Tensor(classes), reduction="mean", ignore_index=classes[0], label_smoothing=0.3)) + + # with batch. somehow this does not match torch + classes = np.random.randint(0, 10, (3,12), dtype=np.int32).tolist() + helper_test_op([(3,12,10)], lambda x: torch.nn.CrossEntropyLoss()(x.permute(0,2,1), torch.tensor(classes)), + lambda x: x.sparse_categorical_crossentropy(Tensor(classes))) + + def test_sparse_categorical_crossentropy_reductions(self): + for r in ("mean", "sum", "none"): + classes = np.random.randint(0, 10, (12,), dtype=np.int32).tolist() + helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(reduction=r)(x, torch.tensor(classes)), + lambda x: x.sparse_categorical_crossentropy(Tensor(classes), reduction=r)) + + def test_sparse_categorical_crossentropy_ignore_index(self): + classes = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] + for i in (-1, 0, 3): + helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(ignore_index=i)(x, torch.tensor(classes)), + lambda x: x.sparse_categorical_crossentropy(Tensor(classes), ignore_index=i)) + + def test_sparse_categorical_crossentropy_label_smoothing(self): + for s in (0.3, 0.9): + classes = np.random.randint(0, 10, (12,), dtype=np.int32).tolist() + helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(label_smoothing=s)(x, torch.tensor(classes)), + lambda x: x.sparse_categorical_crossentropy(Tensor(classes), label_smoothing=s)) + def test_nll_loss(self): target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist() helper_test_op([(32,10)],