chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,220 @@
|
||||
// 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) {}
|
||||
Reference in New Issue
Block a user