wehub-resource-sync
221 lines
8.2 KiB
C++
221 lines
8.2 KiB
C++
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include <math.h>
|
|
|
|
#include <iterator>
|
|
#include <random>
|
|
#include <set>
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/math/sampler.h"
|
|
#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace phi {
|
|
namespace sr {
|
|
|
|
using Sampler = phi::math::Sampler;
|
|
|
|
template <typename T, typename Context>
|
|
void NCEGradKernel(const Context &dev_ctx,
|
|
const DenseTensor &input_in,
|
|
const DenseTensor &label_in,
|
|
const optional<DenseTensor> &bias_in,
|
|
const DenseTensor &weight_in,
|
|
const DenseTensor &sample_logits_in,
|
|
const DenseTensor &sample_labels_in,
|
|
const optional<DenseTensor> &sample_weight_in,
|
|
const optional<DenseTensor> &custom_dist_probs,
|
|
const optional<DenseTensor> &custom_dist_alias,
|
|
const optional<DenseTensor> &custom_dist_alias_probs,
|
|
const DenseTensor &cost_grad,
|
|
int num_total_classes,
|
|
const std::vector<int> &custom_neg_classes,
|
|
int num_neg_samples,
|
|
int sampler_in,
|
|
int seed,
|
|
bool is_sparse,
|
|
bool remote_prefetch,
|
|
bool is_test,
|
|
DenseTensor *input_grad,
|
|
DenseTensor *bias_grad,
|
|
SelectedRows *weight_grad) {
|
|
auto d_out = &cost_grad;
|
|
const T *d_out_data = d_out->data<T>();
|
|
auto label = &label_in;
|
|
auto sample_out = &sample_logits_in;
|
|
const T *sample_out_data = sample_out->data<T>();
|
|
auto sample_labels = &sample_labels_in;
|
|
const int64_t *sample_labels_data = sample_labels->data<int64_t>();
|
|
auto sample_weight = sample_weight_in.get_ptr();
|
|
const T *sample_weight_data = nullptr;
|
|
if (sample_weight != nullptr) {
|
|
sample_weight_data = sample_weight->data<T>();
|
|
}
|
|
int num_true_class = 1;
|
|
if (label != nullptr) {
|
|
num_true_class = label->dims()[1];
|
|
}
|
|
|
|
int sampler_type = sampler_in;
|
|
Sampler *sampler;
|
|
switch (sampler_type) {
|
|
case 0: {
|
|
sampler = new phi::math::UniformSampler(num_total_classes - 1, seed);
|
|
break;
|
|
}
|
|
case 1: {
|
|
sampler = new phi::math::LogUniformSampler(num_total_classes - 1, seed);
|
|
break;
|
|
}
|
|
case 2: {
|
|
auto dist_probs = custom_dist_probs.get_ptr();
|
|
auto dist_alias = custom_dist_alias.get_ptr();
|
|
auto dist_alias_probs = custom_dist_alias_probs.get_ptr();
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
dist_probs->numel(),
|
|
num_total_classes,
|
|
common::errors::InvalidArgument(
|
|
"ShapeError: The number of elements in Input(CustomDistProbs) "
|
|
"should be equal to the number of total classes. But Received: "
|
|
"Input(CustomDistProbs).numel() = %d, Attr(num_total_classes) "
|
|
"= %d.",
|
|
dist_probs->numel(),
|
|
num_total_classes));
|
|
PADDLE_ENFORCE_EQ(
|
|
dist_alias->numel(),
|
|
num_total_classes,
|
|
common::errors::InvalidArgument(
|
|
"ShapeError: The number of elements in Input(CustomDistAlias) "
|
|
"should be equal to the number of total classes. But Received: "
|
|
"Input(CustomDistAlias).numel() = %d, Attr(num_total_classes) "
|
|
"= %d.",
|
|
dist_alias->numel(),
|
|
num_total_classes));
|
|
PADDLE_ENFORCE_EQ(
|
|
dist_alias_probs->numel(),
|
|
num_total_classes,
|
|
common::errors::InvalidArgument(
|
|
"ShapeError: The number of elements in "
|
|
"Input(CustomDistAliasProbs) "
|
|
"should be equal to the number of total classes. But Received: "
|
|
"Input(CustomDistAliasProbs).numel() = %d, "
|
|
"Attr(num_total_classes) = %d.",
|
|
dist_alias_probs->numel(),
|
|
num_total_classes));
|
|
|
|
const float *probs_data = dist_probs->data<float>();
|
|
const int *alias_data = dist_alias->data<int>();
|
|
const float *alias_probs_data = dist_alias_probs->data<float>();
|
|
sampler = new phi::math::CustomSampler(num_total_classes - 1,
|
|
probs_data,
|
|
alias_data,
|
|
alias_probs_data,
|
|
seed);
|
|
break;
|
|
}
|
|
default: {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"Unsupported SamplerType. SamplerType should be 0: Uniform, "
|
|
"1: LogUniform or 2: CustomDist. Received SamplerType: %d",
|
|
sampler_type));
|
|
}
|
|
}
|
|
|
|
// T b = 1. / num_total_classes * num_neg_samples;
|
|
DenseTensor sample_grad; // tmp tensor
|
|
sample_grad.Resize(sample_labels->dims());
|
|
T *sample_grad_data = dev_ctx.template Alloc<T>(&sample_grad);
|
|
|
|
// backward cost
|
|
for (int64_t i = 0; i < sample_labels->numel(); ++i) {
|
|
int64_t label_idx = i % sample_labels->dims()[1];
|
|
int64_t sample_idx = i / sample_labels->dims()[1];
|
|
float b = sampler->Probability(sample_labels_data[i]) * num_neg_samples;
|
|
T o = sample_out_data[i];
|
|
T w = sample_weight == nullptr ? 1 : sample_weight_data[sample_idx];
|
|
sample_grad_data[i] = label_idx < num_true_class
|
|
? w * (b / (o + b)) * (o - 1)
|
|
: w * (o * (1 - o) / (o + b));
|
|
sample_grad_data[i] *= d_out_data[sample_idx];
|
|
}
|
|
|
|
// get d_bias
|
|
auto d_bias = bias_grad;
|
|
if (d_bias != nullptr) {
|
|
T *d_bias_data = dev_ctx.template Alloc<T>(d_bias);
|
|
std::fill(d_bias_data, d_bias_data + d_bias->numel(), 0.0);
|
|
for (int64_t i = 0; i < sample_labels->numel(); ++i) {
|
|
d_bias_data[sample_labels_data[i]] += sample_grad_data[i];
|
|
}
|
|
}
|
|
|
|
if (!is_sparse) {
|
|
PADDLE_THROW(
|
|
common::errors::InvalidArgument("The parameter weight_grad of a NCE_OP "
|
|
"must be DenseTensor"));
|
|
} else {
|
|
std::vector<int64_t> labels;
|
|
for (int64_t i = 0; i < sample_labels->numel(); ++i) {
|
|
labels.push_back(sample_labels_data[i]);
|
|
}
|
|
std::set<T> st(labels.begin(), labels.end());
|
|
labels.assign(st.begin(), st.end());
|
|
|
|
DDim table_dim = weight_in.dims();
|
|
|
|
auto d_w = weight_grad;
|
|
|
|
d_w->set_rows(labels);
|
|
d_w->set_height(table_dim[0]);
|
|
|
|
auto *d_table_value = d_w->mutable_value();
|
|
d_table_value->Resize({static_cast<int64_t>(labels.size()), table_dim[1]});
|
|
auto d_w_data = dev_ctx.template Alloc<T>(d_table_value);
|
|
std::fill(d_w_data, d_w_data + d_table_value->numel(), 0.0);
|
|
|
|
auto d_w_matrix = EigenMatrix<T>::From(*d_table_value);
|
|
auto x_matrix = EigenMatrix<T>::From(input_in);
|
|
for (int64_t i = 0; i < sample_labels->numel(); ++i) {
|
|
d_w_matrix.chip(d_w->Index(sample_labels_data[i]), 0) +=
|
|
x_matrix.chip(static_cast<int>(i / sample_labels->dims()[1]), 0) *
|
|
sample_grad_data[i];
|
|
}
|
|
}
|
|
|
|
// get d_x
|
|
auto d_x = input_grad;
|
|
if (d_x != nullptr) {
|
|
auto *d_x_data = dev_ctx.template Alloc<T>(d_x);
|
|
std::fill(d_x_data, d_x_data + d_x->numel(), 0.0);
|
|
auto d_x_matrix = EigenMatrix<T>::From(*d_x);
|
|
auto w_matrix = EigenMatrix<T>::From(weight_in);
|
|
for (int64_t i = 0; i < sample_labels->numel(); ++i) {
|
|
d_x_matrix.chip(static_cast<int>(i / sample_labels->dims()[1]), 0) +=
|
|
w_matrix.chip(sample_labels_data[i], 0) * sample_grad_data[i];
|
|
}
|
|
}
|
|
|
|
delete sampler;
|
|
}
|
|
|
|
} // namespace sr
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
nce_sr_grad, CPU, ALL_LAYOUT, phi::sr::NCEGradKernel, float, double) {}
|