Method Compute

oneflow/user/kernels/nll_kernel.cpp:132–162

  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }

  void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state,
               const user_op::OpKernelCache* cache) const override {
    const auto* target = ctx->Tensor4ArgNameAndIndex("target", 0);
    const auto* out_grad = ctx->Tensor4ArgNameAndIndex("out_grad", 0);
    auto* in_grad = ctx->Tensor4ArgNameAndIndex("in_grad", 0);

    const int64_t N = target->shape_view().elem_cnt();
    const int64_t C = in_grad->shape_view().At(in_grad->shape_view().NumAxes() - 1);
    CHECK_LE(N, std::numeric_limits<int32_t>::max())
        << "Expected batch size not exceed int32 numeric limits";

    K class_start = 0;
    if (cache) {
      const auto* spec_cache = dynamic_cast<const NLLKernelCache*>(cache);
      CHECK_NOTNULL(spec_cache);
      CHECK_EQ(spec_cache->num_classes(), C) << ctx->op_name() << ": expected num_classes " << C
                                             << ", got " << spec_cache->num_classes();
      class_start = spec_cache->class_start();
    }

    const K ignore_index = static_cast<K>(ctx->Attr<int64_t>("ignore_index"));

    const T* weight_dptr = nullptr;
    if (ctx->has_input("weight", 0)) {
      weight_dptr = CHECK_NOTNULL(ctx->Tensor4ArgNameAndIndex("weight", 0))->dptr<T>();
    }

    NLLKernelUtil<device_type, T, K>::Backward(
        ctx->stream(), static_cast<int32_t>(N), static_cast<K>(C), class_start, ignore_index,
        out_grad->dptr<T>(), target->dptr<K>(), weight_dptr, in_grad->mut_dptr<T>());
  }
};

#define REGISTER_NLL_KERNELS(device, dtype, ltype)                                            \