diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 844260a64b..643a87fca2 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -163,6 +163,22 @@ class TestTensorUOpLoss(unittest.TestCase): t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32) self.assertIs(_strip_unique(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop), _strip_unique(t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0))) + def test_nll_loss(self): + t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32) + self.assertIs(_strip_unique(t.nll_loss(Y).uop), _strip_unique(t.uop.nll_loss(Y.uop))) + def test_nll_loss_weight(self): + t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float() + self.assertIs(_strip_unique(t.nll_loss(Y, weight=w).uop), _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop))) + def test_nll_loss_ignore_index(self): + t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32) + self.assertIs(_strip_unique(t.nll_loss(Y, ignore_index=1).uop), _strip_unique(t.uop.nll_loss(Y.uop, ignore_index=1))) + def test_nll_loss_none_reduction(self): + t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32) + self.assertIs(_strip_unique(t.nll_loss(Y, reduction="none").uop), _strip_unique(t.uop.nll_loss(Y.uop, reduction="none"))) + def test_nll_loss_weight_ignore_index(self): + t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float() + self.assertIs(_strip_unique(t.nll_loss(Y, weight=w, ignore_index=1).uop), + _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1))) class TestTensorUOpScatter(unittest.TestCase): def test_scatter(self): diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 0e408bf501..b6c3265ce1 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -1318,6 +1318,30 @@ class OpMixin(ElementwiseMixin, ReduceMixin): 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:Self, weight:Self|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Self: + """ + Computes the negative log likelihood loss between log-probabilities and target labels. + + NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities. + + See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[-1, 2, -3], [1, -2, 3]]) + Y = Tensor([1, 2]) + print(t.log_softmax().nll_loss(Y).item()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[-1, 2, -3], [1, -2, 3]]) + Y = Tensor([1, 2]) + print(t.log_softmax().nll_loss(Y, reduction='none').numpy()) + ``` + """ + weight = Y.ones_like() if weight is None else weight.gather(0, Y.flatten()).reshape(Y.shape) + masked_weight = weight if ignore_index is None else weight * Y.ne(ignore_index) + nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight + return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction) + # ***** matrix ops ***** def newton_schulz(self, steps:int, params:tuple[int, ...], eps:float=1.0e-7) -> Self: diff --git a/tinygrad/nn/onnx.py b/tinygrad/nn/onnx.py index 3a1d149bb2..ecc8de6ee3 100644 --- a/tinygrad/nn/onnx.py +++ b/tinygrad/nn/onnx.py @@ -2,7 +2,8 @@ from typing import Any, Sequence, cast, Literal, NamedTuple, Generator import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum from tinygrad.nn.state import TensorIO -from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr +from tinygrad.tensor import Tensor, _broadcast_shape +from tinygrad.mixin import ReductionStr from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT from tinygrad.device import is_dtype_supported, Device diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 861fe928a9..211190285c 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -10,7 +10,7 @@ from tinygrad.helpers import argfix, flatten, prod, all_int, round_up, getenv, a from tinygrad.helpers import resolve_pool_pads, IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey, cpu_profile from tinygrad.helpers import suppress_finalizing, disable_gc from tinygrad.gradient import compute_gradient -from tinygrad.mixin import OpMixin, ReductionStr +from tinygrad.mixin import OpMixin from tinygrad.uop.ops import smax, UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape from tinygrad.schedule import ExecItem, create_linear_with_vars, linear_to_schedule from tinygrad.device import Buffer, canonicalize_device @@ -1533,30 +1533,6 @@ class Tensor(OpMixin): qk = qk + attn_mask return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value - def nll_loss(self, Y:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Tensor: - """ - Computes the negative log likelihood loss between log-probabilities and target labels. - - NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities. - - See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[-1, 2, -3], [1, -2, 3]]) - Y = Tensor([1, 2]) - print(t.log_softmax().nll_loss(Y).item()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[-1, 2, -3], [1, -2, 3]]) - Y = Tensor([1, 2]) - print(t.log_softmax().nll_loss(Y, reduction='none').numpy()) - ``` - """ - weight = Y.ones_like(requires_grad=False) if weight is None else weight[Y] - masked_weight = weight if ignore_index is None else weight * (Y != ignore_index) - nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight - return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction) - def qr(self) -> tuple[Tensor, Tensor]: assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}" b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1])