| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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. |
| */ |
| |
| /*! |
| * \file pdf_op.h |
| * \brief Operators for computing the pdf of random distributions. |
| */ |
| #ifndef MXNET_OPERATOR_RANDOM_PDF_OP_H_ |
| #define MXNET_OPERATOR_RANDOM_PDF_OP_H_ |
| |
| #include <mxnet/operator_util.h> |
| #include <vector> |
| #include <algorithm> |
| #include "../mshadow_op.h" |
| #include "../mxnet_op.h" |
| #include "../operator_common.h" |
| #include "../elemwise_op_common.h" |
| #include "../special_functions-inl.h" |
| #include "../tensor/broadcast_reduce_op.h" |
| |
| namespace mxnet { |
| namespace op { |
| |
| template <typename DType> |
| MSHADOW_XINLINE static DType ceph_psi(DType val) { |
| return special_functions::cephes::psi(val); |
| } |
| template <> |
| MSHADOW_XINLINE mshadow::half::half_t ceph_psi(mshadow::half::half_t val) { |
| return special_functions::cephes::psi<float>(val); |
| } |
| |
| template <bool logpdf> |
| struct PDF_Uniform { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* lower, |
| IType2* upper) { |
| const index_t index(start / sample_size); |
| const DType l(lower[index]), h(upper[index]); |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| // No check whether sample is in the support. |
| out[i] = logpdf ? -DType(log(h - l)) : DType(1.0) / (h - l); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Uniform_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* lower, |
| IType2* upper, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_lower, |
| IType2* grad_upper) { |
| const index_t index(start / sample_size); |
| const DType l(lower[index]), h(upper[index]); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType scaling(grad_out[i] * (logpdf ? DType(1) : out[i])); |
| grad_lower[i] = scaling / (h - l); |
| grad_upper[i] = scaling / (l - h); |
| KERNEL_ASSIGN(grad_sample[i], req, 0); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Normal { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* loc, |
| IType2* scale) { |
| const index_t index(start / sample_size); |
| const DType u(loc[index]), s(scale[index]), sq(s * s); |
| const DType normalizer(sqrt(2.0 * mxnet_op::PI) * s); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType exponent((DType(-0.5) * (x - u) * (x - u)) / (sq)); |
| out[i] = logpdf ? exponent - log(normalizer) : exp(exponent) / normalizer; |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Normal_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* loc, |
| IType2* scale, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_loc, |
| IType2* grad_scale) { |
| const index_t index(start / sample_size); |
| const DType u(loc[index]), s(scale[index]), s_squared(s * s), s_cubed(s_squared * s); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType scaling(grad_out[i] * (logpdf ? DType(1) : out[i])); |
| grad_loc[i] = scaling * (x - u) / s_squared; |
| grad_scale[i] = scaling * ((x - u) * (x - u) - s_squared) / s_cubed; |
| KERNEL_ASSIGN(grad_sample[i], req, scaling * (u - x) / s_squared); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Gamma { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* alpha, |
| IType2* beta) { |
| const index_t index(start / sample_size); |
| const DType a(alpha[index]), b(beta[index]), lgamma_a(lgamma(a)), a_log_b(a * log(b)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType lpdf(a_log_b + (a - 1) * log(x) - b * x - lgamma_a); |
| out[i] = logpdf ? lpdf : DType(exp(lpdf)); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Gamma_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* alpha, |
| IType2* beta, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_alpha, |
| IType2* grad_beta) { |
| const index_t index(start / sample_size); |
| const DType a(alpha[index]), b(beta[index]), log_b(log(b)), ceph_psi_a(ceph_psi(a)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType scaling(grad_out[i] * (logpdf ? DType(1) : out[i])); |
| grad_alpha[i] = scaling * (log_b + log(x) - ceph_psi_a); |
| grad_beta[i] = scaling * (a / b - x); |
| KERNEL_ASSIGN(grad_sample[i], req, scaling * ((a - 1) / x - b)); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Exponential { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* lambda) { |
| const index_t index(start / sample_size); |
| const DType l(lambda[index]), log_l(log(l)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| out[i] = logpdf ? log_l - l * x : l * exp(-l * x); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Exponential_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* lambda, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_lambda) { |
| const index_t index(start / sample_size); |
| const DType l(lambda[index]); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType scaling(grad_out[i] * (logpdf ? DType(1) : out[i])); |
| grad_lambda[i] = scaling * (DType(1) / l - x); |
| KERNEL_ASSIGN(grad_sample[i], req, -scaling * l); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Poisson { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* lambda) { |
| const index_t index(start / sample_size); |
| const DType l(lambda[index]), log_l(log(l)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType lpdf((x * log_l - lgamma(x + 1)) - l); |
| out[i] = logpdf ? lpdf : DType(exp(lpdf)); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Poisson_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* lambda, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_lambda) { |
| const index_t index(start / sample_size); |
| const DType l(lambda[index]); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType scaling(grad_out[i] * (logpdf ? DType(1) : out[i])); |
| grad_lambda[i] = scaling * (x / l - DType(1)); |
| KERNEL_ASSIGN(grad_sample[i], req, 0); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_NegativeBinomial { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* limit, |
| IType2* prob) { |
| const index_t index(start / sample_size); |
| const DType l(limit[index]), p(prob[index]), lgamma_l(lgamma(l)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType x(sample[i]); |
| const DType lpdf((lgamma(x + l) - lgamma(x + 1) - lgamma_l) + l * log(p) + x * log(1 - p)); |
| out[i] = logpdf ? lpdf : DType(exp(lpdf)); |
| } |
| } |
| |
| template <typename DType> |
| MSHADOW_XINLINE static DType LPDF(DType l, DType p, DType x) { |
| // Note that "p" is the failure and not the success probability. |
| return (lgamma(x + l) - lgamma(x + 1) - lgamma(l)) + l * log(p) + x * log(1 - p); |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_NegativeBinomial_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* limit, |
| IType2* prob, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_limit, |
| IType2* grad_prob) { |
| const index_t index(start / sample_size); |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| DType grad_l(0), grad_p(0); |
| LPDF_GRAD(DType(limit[index]), |
| DType(prob[index]), |
| DType(sample[i]), |
| out[i], |
| grad_out[i], |
| &grad_l, |
| &grad_p); |
| grad_limit[i] = grad_l; |
| grad_prob[i] = grad_p; |
| KERNEL_ASSIGN(grad_sample[i], req, 0); |
| } |
| } |
| |
| template <typename DType> |
| MSHADOW_XINLINE static void |
| LPDF_GRAD(DType l, DType p, DType x, DType o, DType grad_o, DType* grad_l, DType* grad_p) { |
| const DType scaling(grad_o * (logpdf ? DType(1) : o)); |
| *grad_l = scaling * ((ceph_psi(x + l) - ceph_psi(l)) + log(p)); |
| *grad_p = scaling * (l / p - x / (1 - p)); |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_GeneralizedNegativeBinomial { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| DType* out, |
| IType1* sample, |
| IType2* mu, |
| IType2* alpha) { |
| const index_t index(start / sample_size); |
| |
| // Reparameterize with limit = 1 / alpha, prob = 1 / (mu * alpha + 1) |
| const DType limit(1.0 / alpha[index]), prob(1.0 / (mu[index] * alpha[index] + 1.0)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const DType lpdf(PDF_NegativeBinomial<logpdf>::LPDF(limit, prob, DType(sample[i]))); |
| out[i] = logpdf ? lpdf : DType(exp(lpdf)); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_GeneralizedNegativeBinomial_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| DType* out, |
| IType1* sample, |
| IType2* mu, |
| IType2* alpha, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_mu, |
| IType2* grad_alpha) { |
| const index_t index(start / sample_size); |
| const DType fmu(mu[index]), falpha(alpha[index]), den(fmu * falpha + 1.0); |
| |
| // Reparameterize with limit = 1 / alpha, prob = 1 / (mu * alpha + 1) |
| const DType limit(1.0 / falpha), prob(1.0 / (fmu * falpha + 1.0)); |
| |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| // Grad returned as d_limit, d_prob |
| DType grad_l(0), grad_p(0); |
| PDF_NegativeBinomial_Grad<logpdf>::LPDF_GRAD( |
| limit, prob, DType(sample[i]), out[i], grad_out[i], &grad_l, &grad_p); |
| grad_mu[i] = -grad_p * falpha / (den * den); |
| grad_alpha[i] = -grad_l / (falpha * falpha) - grad_p * fmu / (den * den); |
| KERNEL_ASSIGN(grad_sample[i], req, 0); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Dirichlet { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| index_t k, |
| DType* out, |
| IType1* sample, |
| IType2* alpha) { |
| const index_t index(start / sample_size); |
| const index_t end = start + length; |
| for (index_t i = start; i < end; ++i) { |
| const IType1* cur_sample = sample + i * k; |
| const IType2* cur_alpha = alpha + index * k; |
| DType sum_alpha(0), sum_lgamma(0), sum_sample(0); |
| for (index_t j = 0; j < k; ++j) { |
| sum_alpha += cur_alpha[j]; |
| sum_lgamma += lgamma(cur_alpha[j]); |
| sum_sample += (cur_alpha[j] - 1) * log(cur_sample[j]); |
| } |
| DType lpdf(sum_sample + (lgamma(sum_alpha) - sum_lgamma)); |
| out[i] = logpdf ? lpdf : DType(exp(lpdf)); |
| } |
| } |
| }; |
| |
| template <bool logpdf> |
| struct PDF_Dirichlet_Grad { |
| template <typename DType, typename IType1, typename IType2> |
| MSHADOW_XINLINE static void Map(index_t start, |
| index_t length, |
| index_t sample_size, |
| OpReqType req, |
| index_t k, |
| DType* out, |
| IType1* sample, |
| IType2* alpha, |
| DType* grad_out, |
| IType1* grad_sample, |
| IType2* grad_alpha) { |
| const index_t index(start / sample_size); |
| const index_t end = start + length; |
| |
| for (index_t i = start; i < end; ++i) { |
| // Digamma function |
| const IType1* cur_sample = sample + i * k; |
| const IType2* cur_alpha = alpha + index * k; |
| |
| const DType scaling(grad_out[i] * (logpdf ? DType(1) : out[i])); |
| DType sum_alpha(0); |
| for (index_t j = 0; j < k; ++j) { |
| sum_alpha += cur_alpha[j]; |
| } |
| const DType psi_sum(ceph_psi(sum_alpha)); |
| |
| for (index_t j = 0; j < k; ++j) { |
| size_t grad_alpha_index = i % sample_size + sample_size * (j + k * index); |
| size_t grad_sample_index = i * k + j; |
| |
| // order grad_alpha differently to allow efficient reduction at the end. |
| grad_alpha[grad_alpha_index] = |
| scaling * (log(cur_sample[j]) + (psi_sum - ceph_psi(cur_alpha[j]))); |
| KERNEL_ASSIGN( |
| grad_sample[grad_sample_index], req, scaling * (cur_alpha[j] - 1) / cur_sample[j]); |
| } |
| } |
| } |
| }; |
| |
| struct PdfParam : public dmlc::Parameter<PdfParam> { |
| bool is_log; |
| DMLC_DECLARE_PARAMETER(PdfParam) { |
| DMLC_DECLARE_FIELD(is_log).set_default(false).describe( |
| "If set, compute the density of the log-probability instead of the probability."); |
| } |
| }; |
| |
| template <bool vparm = false> |
| inline bool PdfOpShape(const nnvm::NodeAttrs& attrs, |
| std::vector<TShape>* in_attrs, |
| std::vector<TShape>* out_attrs) { |
| CHECK_GT(in_attrs->size(), 1) << "pdf operator takes at least 2 arguments (" << in_attrs->size() |
| << " given)"; |
| CHECK_EQ(out_attrs->size(), 1); |
| // All inputs must be defined in order to infer output shape. |
| if (std::all_of( |
| (*in_attrs).begin(), (*in_attrs).end(), [](const TShape& s) { return s.ndim() > 0; })) { |
| // Tensors of distribution parameters must have same shape. |
| for (size_t i = 2; i < in_attrs->size(); ++i) { |
| SHAPE_ASSIGN_CHECK(*in_attrs, i, (*in_attrs)[i - 1]); |
| } |
| // Tensors of distribution parameters must match leftmost subshape of samples. |
| CHECK_LE((*in_attrs)[1].ndim(), (*in_attrs)[0].ndim()) |
| << "dimension of input samples (" << (*in_attrs)[0].ndim() |
| << ") must be at least dimension of distribution parameters (" << (*in_attrs)[1].ndim() |
| << ")"; |
| TShape tshape((*in_attrs)[0].begin(), (*in_attrs)[0].begin() + (*in_attrs)[1].ndim()); |
| if (vparm) { |
| *(tshape.end() - 1) = *((*in_attrs)[0].end() - 1); |
| } |
| for (size_t i = 1; i < in_attrs->size(); ++i) { |
| SHAPE_ASSIGN_CHECK(*in_attrs, i, tshape); |
| } |
| // Output shape must equal input tensor of samples except for last dimension if we are |
| // dealing with samples that are itself vectors. Be aware of the special case where we |
| // are dealing with a single vector sample. |
| if (vparm && ((*in_attrs)[0].ndim() == 1)) { |
| // Special case where we are dealing with a single vector sample. |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::Shape1(1)); |
| } else { |
| TShape oshape((*in_attrs)[0].begin(), (*in_attrs)[0].end() - (vparm ? 1 : 0)); |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, oshape); |
| } |
| return true; |
| } |
| return false; |
| } |
| |
| template <typename OP> |
| struct LaunchExWrapper { |
| template <typename... Args> |
| MSHADOW_XINLINE static void Map(const index_t start, |
| const index_t length, |
| const index_t sample_size, |
| Args... args) { |
| // Apply the operator to the sample in strides of sample_size, so that |
| // the operators can assume that their distribution parameters are constant. |
| index_t i = start; |
| |
| // Get aligned |
| const index_t align_step = sample_size - (i % sample_size); |
| const index_t first_stride = length > align_step ? align_step : length; |
| OP::Map(i, first_stride, sample_size, args...); |
| i += first_stride; |
| |
| const index_t end = start + length - sample_size; |
| for (; i < end; i += sample_size) { |
| OP::Map(i, sample_size, sample_size, args...); |
| } |
| |
| // Last stride might not be aligned either |
| const index_t last_stride = start + length - i; |
| if (last_stride > 0) { // Don't overstep even if length <= sample_size |
| OP::Map(i, last_stride, sample_size, args...); |
| } |
| } |
| }; |
| |
| template <typename xpu, typename DType, typename pdf, int pnum, bool vparm = false> |
| struct PdfCaller; |
| |
| template <typename xpu, typename DType, typename pdf> |
| struct PdfCaller<xpu, DType, pdf, 1, false> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| mshadow::Stream<xpu>* s) { |
| CHECK_EQ(inputs[0].Size() % inputs[1].Size(), 0); |
| CHECK_EQ(inputs[0].Size() % outputs[0].Size(), 0); |
| index_t num_samples(inputs[0].Size() / inputs[1].Size()); |
| mxnet_op::Kernel<LaunchExWrapper<pdf>, xpu>::LaunchEx(s, |
| outputs[0].Size(), |
| num_samples, |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| inputs[1].dptr<DType>()); |
| } |
| }; |
| |
| template <typename xpu, typename DType, typename pdf> |
| struct PdfCaller<xpu, DType, pdf, 1, true> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| mshadow::Stream<xpu>* s) { |
| CHECK_EQ(inputs[0].Size() % inputs[1].Size(), 0); |
| CHECK_EQ(inputs[0].Size() % outputs[0].Size(), 0); |
| index_t num_samples(inputs[0].Size() / inputs[1].Size()); |
| index_t sample_size(inputs[0].Size() / outputs[0].Size()); |
| |
| // Covers distributions parametrized by a vector of parameters (Dirichlet distribution). |
| mxnet_op::Kernel<LaunchExWrapper<pdf>, xpu>::LaunchEx(s, |
| outputs[0].Size(), |
| num_samples, |
| sample_size, |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| inputs[1].dptr<DType>()); |
| } |
| }; |
| |
| template <typename xpu, typename DType, typename pdf> |
| struct PdfCaller<xpu, DType, pdf, 2, false> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<TBlob>& outputs, |
| mshadow::Stream<xpu>* s) { |
| CHECK_EQ(inputs[0].Size() % inputs[1].Size(), 0); |
| CHECK_EQ(inputs[0].Size(), outputs[0].Size()); |
| index_t num_samples(inputs[0].Size() / inputs[1].Size()); |
| mxnet_op::Kernel<LaunchExWrapper<pdf>, xpu>::LaunchEx(s, |
| outputs[0].Size(), |
| num_samples, |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| inputs[1].dptr<DType>(), |
| inputs[2].dptr<DType>()); |
| } |
| }; |
| |
| template <typename xpu, template <bool> class pdf, int pnum, bool vparm> |
| void PdfOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| CHECK_NE(req[0], kAddTo); |
| CHECK_EQ(inputs.size(), pnum + 1); |
| CHECK_EQ(outputs.size(), 1); |
| |
| // Skip kernel launch for zero-size tensors |
| if (inputs[1].shape_.Size() == 0U || outputs[0].Size() == 0U) { |
| return; |
| } |
| |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| const PdfParam& param = nnvm::get<PdfParam>(attrs.parsed); |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| if (param.is_log) { |
| PdfCaller<xpu, DType, pdf<true>, pnum, vparm>::op(inputs, outputs, s); |
| } else { |
| PdfCaller<xpu, DType, pdf<false>, pnum, vparm>::op(inputs, outputs, s); |
| } |
| }); |
| } |
| |
| template <typename xpu, typename DType, typename pdfgrad, int pnum, int vparm = false> |
| struct PdfGradCaller; |
| |
| template <typename xpu, typename DType, typename pdfgrad> |
| struct PdfGradCaller<xpu, DType, pdfgrad, 1, false> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& grads, |
| mshadow::Stream<xpu>* s) { |
| index_t num_samples(inputs[1].Size() / inputs[2].Size()); |
| mxnet_op::Kernel<LaunchExWrapper<pdfgrad>, xpu>::LaunchEx(s, |
| inputs[0].Size(), |
| num_samples, |
| req[0], |
| inputs[3].dptr<DType>(), |
| inputs[1].dptr<DType>(), |
| inputs[2].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| grads[0].dptr<DType>(), |
| grads[1].dptr<DType>()); |
| } |
| }; |
| |
| template <typename xpu, typename DType, typename pdfgrad> |
| struct PdfGradCaller<xpu, DType, pdfgrad, 1, true> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& grads, |
| mshadow::Stream<xpu>* s) { |
| index_t num_samples(inputs[1].Size() / inputs[2].Size()); |
| index_t sample_size(inputs[1].Size() / inputs[0].Size()); |
| mxnet_op::Kernel<LaunchExWrapper<pdfgrad>, xpu>::LaunchEx(s, |
| inputs[0].Size(), |
| num_samples, |
| req[0], |
| sample_size, |
| inputs[3].dptr<DType>(), |
| inputs[1].dptr<DType>(), |
| inputs[2].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| grads[0].dptr<DType>(), |
| grads[1].dptr<DType>()); |
| } |
| }; |
| |
| template <typename xpu, typename DType, typename pdfgrad> |
| struct PdfGradCaller<xpu, DType, pdfgrad, 2, false> { |
| static void op(const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& grads, |
| mshadow::Stream<xpu>* s) { |
| index_t num_samples(inputs[1].Size() / inputs[2].Size()); |
| mxnet_op::Kernel<LaunchExWrapper<pdfgrad>, xpu>::LaunchEx(s, |
| inputs[0].Size(), |
| num_samples, |
| req[0], |
| inputs[4].dptr<DType>(), |
| inputs[1].dptr<DType>(), |
| inputs[2].dptr<DType>(), |
| inputs[3].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| grads[0].dptr<DType>(), |
| grads[1].dptr<DType>(), |
| grads[2].dptr<DType>()); |
| } |
| }; |
| |
| template <typename xpu, template <bool> class pdfgrad, int pnum, bool vparm> |
| void PdfOpBackward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| CHECK_EQ(inputs.size(), pnum + 3); |
| CHECK_EQ(outputs.size(), pnum + 1); |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| const PdfParam& param = nnvm::get<PdfParam>(attrs.parsed); |
| const size_t N(outputs[1].Size()); |
| const TShape src_shape(Shape2(N, outputs[0].Size() / N)), dst_shape(Shape2(N, 1)); |
| const size_t red_work_size(broadcast::ReduceWorkspaceSize(s, dst_shape, kAddTo, src_shape)); |
| #if !defined(__CUDACC__) |
| // Inputs to PdfOpBackward: grad, samples, parm1, parm2, pdf. |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| const size_t tmp_size(outputs[0].Size() * pnum * sizeof(DType) + red_work_size); |
| Tensor<xpu, 1, char> tmp_space = |
| ctx.requested[0].get_space_typed<xpu, 1, char>(Shape1(tmp_size), s); |
| std::vector<TBlob> grads = {outputs[0]}; |
| grads.push_back(TBlob(tmp_space.dptr_, |
| outputs[0].shape_, |
| outputs[1].dev_mask(), |
| outputs[1].type_flag_, |
| outputs[1].dev_id())); |
| if (pnum == 2) { |
| grads.push_back(TBlob(tmp_space.dptr_ + outputs[0].Size() * sizeof(DType), |
| outputs[0].shape_, |
| outputs[2].dev_mask(), |
| outputs[2].type_flag_, |
| outputs[2].dev_id())); |
| } |
| if (param.is_log) { |
| PdfGradCaller<xpu, DType, pdfgrad<true>, pnum, vparm>::op(inputs, req, grads, s); |
| } else { |
| PdfGradCaller<xpu, DType, pdfgrad<false>, pnum, vparm>::op(inputs, req, grads, s); |
| } |
| Tensor<xpu, 1, char> red_work( |
| tmp_space.dptr_ + pnum * outputs[0].Size() * sizeof(DType), Shape1(red_work_size), s); |
| broadcast::Reduce<red::sum, 2, DType, op::mshadow_op::identity>( |
| s, outputs[1].reshape(dst_shape), req[1], red_work, grads[1].reshape(src_shape)); |
| if (pnum == 2) { |
| broadcast::Reduce<red::sum, 2, DType, op::mshadow_op::identity>( |
| s, outputs[2].reshape(dst_shape), req[2], red_work, grads[2].reshape(src_shape)); |
| } |
| }); |
| #else |
| // Inputs to PdfOpBackward: grad, samples, parm1, parm2, pdf. |
| MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| const size_t tmp_size(outputs[0].Size() * pnum * sizeof(DType) + red_work_size); |
| Tensor<xpu, 1, char> tmp_space = |
| ctx.requested[0].get_space_typed<xpu, 1, char>(Shape1(tmp_size), s); |
| std::vector<TBlob> grads = {outputs[0]}; |
| grads.push_back(TBlob(tmp_space.dptr_, |
| outputs[0].shape_, |
| outputs[1].dev_mask(), |
| outputs[1].type_flag_, |
| outputs[1].dev_id())); |
| if (pnum == 2) { |
| grads.push_back(TBlob(tmp_space.dptr_ + outputs[0].Size() * sizeof(DType), |
| outputs[0].shape_, |
| outputs[2].dev_mask(), |
| outputs[2].type_flag_, |
| outputs[2].dev_id())); |
| } |
| if (param.is_log) { |
| PdfGradCaller<xpu, DType, pdfgrad<true>, pnum, vparm>::op(inputs, req, grads, s); |
| } else { |
| PdfGradCaller<xpu, DType, pdfgrad<false>, pnum, vparm>::op(inputs, req, grads, s); |
| } |
| Tensor<xpu, 1, char> red_work( |
| tmp_space.dptr_ + pnum * outputs[0].Size() * sizeof(DType), Shape1(red_work_size), s); |
| broadcast::RTCReduce(ctx, |
| outputs[1].reshape(dst_shape), |
| req[1], |
| red_work, |
| grads[1].reshape(src_shape), |
| "red::sum{}", |
| 2, |
| "identity"); |
| if (pnum == 2) { |
| broadcast::RTCReduce(ctx, |
| outputs[2].reshape(dst_shape), |
| req[2], |
| red_work, |
| grads[2].reshape(src_shape), |
| "red::sum{}", |
| 2, |
| "identity"); |
| } |
| }); |
| |
| #endif |
| } |
| |
| } // namespace op |
| } // namespace mxnet |
| |
| #endif // MXNET_OPERATOR_RANDOM_PDF_OP_H_ |
