Where can I find cross_entropy_loss implementation

hwijeen · February 9, 2022, 1:55am

Hi, I would like to see the implementation of cross entropy loss. So far, I learned that, torch.nn.funcional.py calls torch._C.nn.cross_entropy_loss but I am having trouble finding the C implementation.

From the releated issue(Where does `torch._C` come from?), I realize that some functions are implemented in C and torch._C is the implementation. The Issue contains links to the original implementations, but it from the outdated pytorch code.

Can someone point to the exact location of cross entropy loss implementation (both CPU and GPU)?
If possible, can someone kindly explain how one can track down to such point from python package(e.g. from torch.nn.functional.py to the C implementation, where the C code is binded with python, etc…)

Many thanks!

eqy · February 9, 2022, 2:07am

Cross-entropy loss is implemented with a combination of existing operations, such as logsoftmax and nll_loss:

github.com

pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L568

      
        
                    break;
                  case Reduction::None:
                    ret = smooth_loss;
                    break;
                  default:
                    TORCH_CHECK(false, "Invalid reduction type encountered in cross_entropy: ", reduction);
                }
                return (1 - label_smoothing) * nllloss + ret * (label_smoothing / n_classes);
            }
            
            
Tensor cross_entropy_loss(
                const Tensor& self,
                const Tensor& target,
                const c10::optional<Tensor>& weight,
                int64_t reduction,
                int64_t ignore_index,
                double label_smoothing) {
              Tensor ret;
              if (self.sizes() == target.sizes()) {
                // Assume soft targets when input and target shapes are the same
                TORCH_CHECK(at::isFloatingType(target.scalar_type()),

These individual kernels will have CPU and GPU implementations in ATen, e.g.,
CPU:

github.com

pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L274

      
        
            
            
  if (reduction == Reduction::Mean) {
                output_val /= total_weight_val;
              }
            
            
  // write result to output tensors
              *output.data_ptr<scalar_t>() = output_val;
              *total_weight_data = total_weight_val;
            }
            
            
void nll_loss_forward_out_cpu_template(
                const Tensor& output,
                const Tensor& total_weight,
                const Tensor& input,
                const Tensor& target,
                const Tensor& weight,
                int64_t reduction,
                int64_t ignore_index) {
              AT_DISPATCH_FLOATING_TYPES_AND(
                  ScalarType::BFloat16, input.scalar_type(), "nll_loss_out_frame", [&] {
                    if (target.scalar_type() == kByte) {

github.com

pytorch/pytorch/blob/2981534f54d49fa3a9755c9b0855e7929c2527f0/aten/src/ATen/native/LossNLL.cpp#L133

      
        
            }
            
            
// Returns the address of the first element of a tensor
            // or nullptr if the tensor is undefined.
            template <typename scalar_t>
            inline scalar_t* optional_data(const Tensor& source) {
              return source.defined() ? source.data_ptr<scalar_t>() : nullptr;
            }
            
            
template <typename scalar_t, typename target_t>
            static void nll_loss_out_frame(
                const Tensor& output,
                const Tensor& total_weight,
                const Tensor& input,
                const Tensor& target,
                const Tensor& weight,
                int64_t reduction,
                int64_t ignore_index) {
              const auto n_dims = input.dim();
              const auto n_classes = input.size(-1);

github.com

pytorch/pytorch/blob/aab67c6dff3819d3c05b328e964f1642227adbe6/aten/src/ATen/native/SoftMax.cpp#L129

      
        
                  "dim must be non-negative and less than input dimensions");
            
            
  set_output(grad.sizes(), grad_input_options);
            }
            }
            
            
namespace native {
            namespace {
            
            
template <typename scalar_t, bool LogSoftMax, bool MaskedSoftMax = false>
            void host_softmax(
                Tensor output,
                const Tensor& input,
                const int64_t dim,
                bool* mask = nullptr) {
              int64_t outer_size = 1;
              int64_t dim_size = input.size(dim);
              int64_t inner_size = 1;
              for (const auto i : c10::irange(dim)) {
                outer_size *= input.size(i);
              }

GPU:

github.com

pytorch/pytorch/blob/aab67c6dff3819d3c05b328e964f1642227adbe6/aten/src/ATen/native/cuda/Loss.cu#L180

      
        
                RECORD_KERNEL_FUNCTION_DTYPE(NAME, _it)                                 \
                switch (_it) {                                                          \
                  AT_PRIVATE_CASE_TYPE_USING_HINT(NAME, at::ScalarType::Byte, uint8_t, index_t, __VA_ARGS__) \
                  AT_PRIVATE_CASE_TYPE_USING_HINT(NAME, at::ScalarType::Long, int64_t, index_t, __VA_ARGS__)\
                  default:                                                              \
                    AT_ERROR(#NAME, " not implemented for '", toString(_it), "'");      \
                }                                                                       \
              }()
            
            
template <typename scalar_t, typename index_t>
            __global__ void nll_loss_forward_no_reduce_cuda_kernel(
                int64_t batch_size,
                PackedTensorAccessor64<scalar_t, 2> input,
                index_t* target,
                scalar_t* output,
                scalar_t* weights,
                int n_classes,
                int ignore_index) {
              CUDA_KERNEL_LOOP(index, batch_size) {
                int cur_target = target[index];
                if (cur_target == ignore_index) {

github.com

pytorch/pytorch/blob/aab67c6dff3819d3c05b328e964f1642227adbe6/aten/src/ATen/native/cuda/SoftMax.cu#L683

      
        
              Epilogue<scalar_t, accscalar_t, outscalar_t> epilogue(sum_k);
            
            
  if (shift == output_shift && shift == grad_output_shift) {
                WriteBpropResultsVectorized<ILP, scalar_t, accscalar_t, outscalar_t, Epilogue>(classes, shift, gradInput, output, gradOutput, epilogue);
              } else {
                WriteBpropResults<ILP, scalar_t, accscalar_t, outscalar_t, Epilogue>(classes, gradInput, output, gradOutput, epilogue);
              }
            }
            
            
template<template<typename, typename, typename> class Epilogue, bool is_log_softmax>
            Tensor host_softmax(const Tensor & input_, const int64_t dim_, const bool half_to_float, const Tensor& output){
              if (half_to_float) {
                TORCH_CHECK(input_.scalar_type() == ScalarType::Half, "conversion is supported for Half type only");
              }
              auto input = input_.contiguous();
              static_assert(std::is_same<acc_type<at::Half, true>, float>::value, "accscalar_t for half should be float");
              if (input.dim() == 0) input = input.view(1);
              int64_t dim = maybe_wrap_dim(dim_, input.dim());
              TORCH_CHECK(dim >=0 && dim < input.dim(), "dim must be non-negative and less than input dimensions");
              int64_t outer_size = 1;
              int64_t dim_size = input.size(dim);