mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
5fefb29a53
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/13856 Codemod generated with clangr shard mode, 25 files per diff, motivation: https://github.com/pytorch/pytorch/pull/12407 Reviewed By: smessmer Differential Revision: D13007310 fbshipit-source-id: 941f064ef8934bb17fbfb706e6ed3db173b5d268
115 lines
3.9 KiB
C++
115 lines
3.9 KiB
C++
#include "caffe2/operators/weighted_multi_sampling_op.h"
|
|
#include "caffe2/utils/math.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
template <class Context>
|
|
bool WeightedMultiSamplingOp<Context>::RunOnDevice() {
|
|
const auto& weight = Input(0);
|
|
CAFFE_ENFORCE_EQ(weight.dim(), 1, "Input should be 1-D vector");
|
|
auto dims = weight.sizes().vec();
|
|
size_t data_size = weight.dim32(0);
|
|
|
|
std::vector<int64_t> indices_sizes;
|
|
auto num_samples = num_samples_;
|
|
if (InputSize() == 2) {
|
|
CAFFE_ENFORCE(
|
|
!OperatorBase::HasArgument("num_samples"),
|
|
"New shape is specified by the input blob, do not pass in "
|
|
"the argument `num_samples`.");
|
|
num_samples = Input(1).numel();
|
|
indices_sizes = Input(1).sizes().vec();
|
|
} else {
|
|
indices_sizes = {num_samples};
|
|
}
|
|
|
|
auto* indices = Output(0, indices_sizes, at::dtype<int>());
|
|
int* indices_data = indices->template mutable_data<int>();
|
|
if (data_size == 0) {
|
|
indices->Resize(0);
|
|
return true;
|
|
}
|
|
|
|
const float* weight_data = weight.template data<float>();
|
|
|
|
for (int i = 0; i < num_samples; ++i) {
|
|
float r;
|
|
math::RandUniform<float, Context>(
|
|
1, 0.0f, weight_data[data_size - 1], &r, &context_);
|
|
auto lb = std::lower_bound(weight_data, weight_data + data_size, r);
|
|
CAFFE_ENFORCE(
|
|
lb != weight_data + data_size, "Cannot find ", r, " in input CDF.");
|
|
indices_data[i] = static_cast<int>(lb - weight_data);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
REGISTER_CPU_OPERATOR(
|
|
WeightedMultiSampling,
|
|
WeightedMultiSamplingOp<CPUContext>);
|
|
|
|
OPERATOR_SCHEMA(WeightedMultiSampling)
|
|
.NumInputs(1, 2)
|
|
.NumOutputs(1)
|
|
.TensorInferenceFunction([](const OperatorDef& def,
|
|
const vector<TensorShape>& in) {
|
|
vector<TensorShape> out(1);
|
|
if (in[0].dims(0) == 0) {
|
|
out[0].set_data_type(TensorProto::INT32);
|
|
out[0].add_dims(0);
|
|
return out;
|
|
}
|
|
|
|
const ArgumentHelper args(def);
|
|
if (args.HasArgument("num_samples")) {
|
|
CAFFE_ENFORCE_EQ(
|
|
in.size(),
|
|
1,
|
|
"New shape must not be specified by the input blob and the "
|
|
"argument `num_samples` at the same time.");
|
|
int num_samples = args.GetSingleArgument<int64_t>("num_samples", 0);
|
|
out[0] =
|
|
CreateTensorShape(vector<int64_t>{num_samples}, TensorProto::INT32);
|
|
return out;
|
|
} else {
|
|
CAFFE_ENFORCE_EQ(
|
|
in.size(),
|
|
2,
|
|
"New shape must be specified by either the input blob or the "
|
|
"argument `num_samples`.");
|
|
std::vector<int64_t> output_dims = GetDimsVector(in[1]);
|
|
out[0] = CreateTensorShape(output_dims, TensorProto::INT32);
|
|
return out;
|
|
}
|
|
})
|
|
.SetDoc(R"DOC(
|
|
The operator performs sampling based on the input sampling weights.
|
|
All weights are cummulative probability thus sorted. The output is
|
|
a 1-D tensor (Tensor). If two inputs are given, the second input
|
|
is used to provide shape of the output sample tensor. Otherwise, we use
|
|
argument `num_samples` to determine the number of samples to generate.
|
|
)DOC")
|
|
.Input(
|
|
0,
|
|
"sampling_cdf",
|
|
"An optional 1-D Tensor."
|
|
"Input cumulative sampling probability (such as [0.2, 0.5, 0.8, 1.5])."
|
|
" All weights must be non-negative numbers. Note that the last value of"
|
|
" CDF is not necessary 1. If the last value is not 1, all values in"
|
|
" sampling_cdf will be scaled by this number.")
|
|
.Input(
|
|
1,
|
|
"shape_tensor (optional)",
|
|
"Tensor whose shape will be applied to output.")
|
|
.Output(
|
|
0,
|
|
"sampled_indexes",
|
|
"The output tensor contains indices sampled from distribution given"
|
|
"by the weight vector in the input tensor"
|
|
"The output is a 1-D Tensor of size determined by argument"
|
|
"`num_samples` or the second input tensor.")
|
|
.Arg("num_samples", "number of samples to sample from the input data");
|
|
|
|
SHOULD_NOT_DO_GRADIENT(WeightedMultiSample);
|
|
} // namespace caffe2
|