src/backends/cpu/nll_loss.cpp - master - vuron/adept

src/backends/cpu/nll_loss.cpp
102 строки4 KB
kolkir
revert compiler compatibility to gcc12
03 мар 2025, 23:52
03 мар 2025, 23:5226a9274
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE "../src/backends/cpu/log_softmax.cpp"
#include <hwy/foreach_target.h>

#include <hwy/highway.h>
#include "arithmetics-inl.hpp"

#if HWY_ONCE

#include <adept/backends/cpu/nll_loss.hpp>
#include <adept/print.hpp>
#include <adept/threading.hpp>
#include <adept/types_dispatch.hpp>

#include <cmath>

namespace adept::cpu {

void nll_loss_fwd_check_inputs(const Tensor& log_props, const Tensor& target) {
  if (log_props.properties().shape.rang() != 2) {
    THROW_ERROR("nll_loss_fwd support only 2d probabilites tensors [N,C], given {}",
                log_props.properties().shape);
  }

  if (target.properties().shape.rang() != 2) {
    THROW_ERROR("nll_loss_fwd support only 2d target tensors [N,1], given {}",
                target.properties().shape);
  }

  if (target.properties().dtype != dtype_t::Int32) {
    THROW_ERROR("nll_loss_fwd support only int32 target tensors, given {}",
                target.properties().dtype);
  }

  if (target.properties().shape.dim(0) != log_props.properties().shape.dim(0)) {
    THROW_ERROR("nll_loss_fwd got different batch number for tensors, ",
                target.properties().shape.dim(0), " != ", log_props.properties().shape.dim(0));
  }

  if (target.properties().shape.dim(1) != 1) {
    THROW_ERROR("nll_loss_fwd got incorrect dim value for target class {}",
                target.properties().shape);
  }
}

Tensor nll_loss_fwd(const Tensor& log_props, const Tensor& target) {
  nll_loss_fwd_check_inputs(log_props, target);

  auto num_batches = log_props.properties().shape.dim(0);
  auto num_classes = log_props.properties().shape.dim(1);

  Shape new_shape({num_batches, 1});
  auto result = Tensor::empty(
      TensorProperties{new_shape, log_props.properties().device, log_props.properties().dtype});

  DISPATCH_FLOATING_TYPES(log_props.properties().dtype, [&]() {
    const scalar_t* prop_data_base = log_props.const_data_ptr<scalar_t>();
    const int32_t* target_data_base = target.const_data_ptr<int32_t>();
    scalar_t* result_data_base = result.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, num_batches, [&](auto begin, auto end) {
      for (index_t row_i = begin; row_i < end; ++row_i) {
        const auto* prop_data = prop_data_base + row_i * num_classes;
        const auto* target_data = target_data_base + row_i;
        auto* result_data = result_data_base + row_i;

        auto class_idx = target_data[0];
        result_data[0] = -prop_data[class_idx];
      }
    });
  });

  return result;
}

Tensor nll_loss_bwd(const Tensor& target, const Tensor& out_grad, index_t num_classes) {
  auto num_batches = target.properties().shape.dim(0);

  Shape new_shape({num_batches, num_classes});
  auto input_grad = Tensor::zero(
      TensorProperties{new_shape, out_grad.properties().device, out_grad.properties().dtype});

  DISPATCH_FLOATING_TYPES(out_grad.properties().dtype, [&]() {
    const scalar_t* out_grad_data_base = out_grad.const_data_ptr<scalar_t>();
    const int32_t* target_data_base = target.const_data_ptr<int32_t>();
    scalar_t* in_grad_data_base = input_grad.mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, num_batches, [&](auto begin, auto end) {
      for (index_t row_i = begin; row_i < end; ++row_i) {
        const auto* out_grad_data = out_grad_data_base + row_i;
        const auto* target_data = target_data_base + row_i;
        auto* in_grad_data = in_grad_data_base + row_i * num_classes;

        auto class_idx = target_data[0];
        in_grad_data[class_idx] = -out_grad_data[0];
      }
    });
  });

  return input_grad;
}

}  // namespace adept::cpu
#endif  // HWY_ONCE