diff --git a/src/base/nll_loss.h b/src/base/nll_loss.h index de59a941d..b5dfc5c56 100644 --- a/src/base/nll_loss.h +++ b/src/base/nll_loss.h @@ -2,6 +2,7 @@ #define INFINI_OPS_BASE_NLL_LOSS_H_ #include +#include #include "operator.h" @@ -9,6 +10,31 @@ namespace infini::ops { class NllLoss : public Operator { public: + NllLoss(const Tensor input, const Tensor target, + const std::optional weight, + const std::optional size_average, const int64_t ignore_index, + const std::optional reduce, const std::string reduction, + Tensor out) + : input_shape_{input.shape()}, + input_strides_{input.strides()}, + input_type_{input.dtype()}, + target_shape_{target.shape()}, + target_strides_{target.strides()}, + target_type_{target.dtype()}, + out_shape_{out.shape()}, + out_strides_{out.strides()}, + out_type_{out.dtype()}, + has_weight_{weight.has_value()}, + weight_shape_{weight ? weight->shape() : Tensor::Shape{}}, + weight_strides_{weight ? weight->strides() : Tensor::Strides{}}, + weight_type_{weight ? weight->dtype() : DataType::kFloat32}, + reduction_{ + ReductionFromPythonArguments(size_average, reduce, reduction)}, + ignore_index_{ignore_index}, + device_index_{out.device().index()} {} + + [[deprecated( + "Use the overload with `ignore_index` before string `reduction`.")]] NllLoss(const Tensor input, const Tensor target, const std::optional weight, const int64_t reduction, const int64_t ignore_index, Tensor out) @@ -29,6 +55,19 @@ class NllLoss : public Operator { ignore_index_{ignore_index}, device_index_{out.device().index()} {} + void operator()(const Tensor input, const Tensor target, + const std::optional weight, + const std::optional size_average, + const int64_t ignore_index, const std::optional reduce, + const std::string reduction, Tensor out) const { + return operator()( + input, target, weight, + ReductionFromPythonArguments(size_average, reduce, reduction), + ignore_index, out); + } + + [[deprecated( + "Use the overload with `ignore_index` before string `reduction`.")]] virtual void operator()(const Tensor input, const Tensor target, const std::optional weight, const int64_t reduction, const int64_t ignore_index, @@ -66,6 +105,40 @@ class NllLoss : public Operator { int64_t ignore_index_{}; int device_index_{0}; + + private: + static int64_t ReductionFromPythonArguments( + const std::optional size_average, const std::optional reduce, + const std::string& reduction) { + if (!size_average.has_value() && !reduce.has_value()) { + return ReductionFromString(reduction); + } + + if (!reduce.value_or(true)) { + return 0; + } + + if (!size_average.value_or(true)) { + return 2; + } + + return 1; + } + + static int64_t ReductionFromString(const std::string& reduction) { + if (reduction == "none") { + return 0; + } + + if (reduction == "mean") { + return 1; + } + + assert(reduction == "sum" && + "`NllLoss` reduction must be `none`, `mean`, or `sum`"); + + return 2; + } }; } // namespace infini::ops diff --git a/tests/test_nll_loss.py b/tests/test_nll_loss.py new file mode 100644 index 000000000..eacc3432b --- /dev/null +++ b/tests/test_nll_loss.py @@ -0,0 +1,165 @@ +import infini.ops +import pytest + +import torch +from tests.utils import ( + Payload, + empty_strided, + get_stream, + rand_strided, + randint_strided, + randn_strided, +) + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + "expected_reduction, has_weight, ignore_index, api_style, size_average, reduce, reduction", + ( + ("mean", False, -100, "keyword", None, None, "mean"), + ("none", False, -100, "full", None, None, "none"), + ("mean", True, -100, "full", None, None, "mean"), + ("sum", True, 2, "legacy", None, None, "sum"), + ("sum", False, -100, "full", None, None, "sum"), + ("none", False, -100, "full", None, False, "sum"), + ("sum", True, -100, "full", False, None, "none"), + ("mean", True, -100, "full", None, True, "sum"), + ), +) +@pytest.mark.parametrize( + "dtype, rtol, atol", + ( + (torch.float32, 1e-5, 1e-5), + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 1e-2, 1e-2), + ), +) +def test_nll_loss( + expected_reduction, + has_weight, + ignore_index, + api_style, + size_average, + reduce, + reduction, + implementation_index, + dtype, + device, + rtol, + atol, +): + batch_size = 7 + num_classes = 5 + input = torch.nn.functional.log_softmax( + randn_strided((batch_size, num_classes), None, dtype=dtype, device=device), + dim=1, + ) + target = randint_strided( + 0, + num_classes, + (batch_size,), + None, + dtype=torch.int64, + device=device, + ) + + if ignore_index >= 0: + target[0] = ignore_index + + weight = None + + if has_weight: + weight = rand_strided((num_classes,), None, dtype=dtype, device=device).add_( + 0.5 + ) + + out_shape = target.shape if expected_reduction == "none" else () + out = empty_strided(out_shape, None, dtype=dtype, device=device) + + return Payload( + lambda *args: _nll_loss( + *args, + reduction=reduction, + ignore_index=ignore_index, + api_style=api_style, + size_average=size_average, + reduce=reduce, + implementation_index=implementation_index, + ), + lambda *args: _torch_nll_loss( + *args, reduction=expected_reduction, ignore_index=ignore_index + ), + (input, target, weight, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _nll_loss( + input, + target, + weight, + out, + *, + reduction, + ignore_index, + api_style, + size_average, + reduce, + implementation_index, +): + kwargs = { + "implementation_index": implementation_index, + "stream": get_stream(input.device), + } + + if api_style == "keyword": + infini.ops.nll_loss( + input=input, + target=target, + weight=weight, + size_average=size_average, + ignore_index=ignore_index, + reduce=reduce, + reduction=reduction, + out=out, + **kwargs, + ) + elif api_style == "legacy": + infini.ops.nll_loss( + input, + target, + weight, + {"none": 0, "mean": 1, "sum": 2}[reduction], + ignore_index, + out, + **kwargs, + ) + else: + infini.ops.nll_loss( + input, + target, + weight, + size_average, + ignore_index, + reduce, + reduction, + out, + **kwargs, + ) + + return out + + +def _torch_nll_loss(input, target, weight, out, *, reduction, ignore_index): + result = torch.nn.functional.nll_loss( + input, + target, + weight=weight, + ignore_index=ignore_index, + reduction=reduction, + ) + out.copy_(result) + + return out