You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2020/01/14 23:57:37 UTC
[GitHub] [incubator-mxnet] blchu opened a new pull request #17311: Added
beamsearch_set_finished Operator
blchu opened a new pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311
## Description ##
Added a new operator, beamsearch_set_finished, for updating the scores of finished beams in beam search more efficiently than using alternate operators.
## Checklist ##
### Essentials ###
Please feel free to remove inapplicable items for your PR.
- [ ] Changes are complete (i.e. I finished coding on this PR)
- [ ] All changes have test coverage:
- Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
- [ ] Code is well-documented:
- For user-facing API changes, API doc string has been updated.
- For new C++ functions in header files, their functionalities and arguments are documented.
- Check the API doc at https://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
- [ ] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
### Changes ###
- [ ] New operator, beamsearch_set_finished, and appropriate unit test
## Comments ##
- Operator will be added to contrib
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] haojin2 commented on a change in pull request
#17311: Added beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
haojin2 commented on a change in pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#discussion_r367028420
##########
File path: tests/python/gpu/test_operator_gpu.py
##########
@@ -1698,6 +1698,54 @@ def test_take_with_type():
'take_a': 'write'},
arg_params=arg_params)
+@with_seed()
Review comment:
one more blank line above
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] haojin2 commented on a change in pull request
#17311: Added beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
haojin2 commented on a change in pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#discussion_r367027893
##########
File path: src/operator/contrib/beamsearch_set_finished-inl.h
##########
@@ -0,0 +1,148 @@
+#include <dmlc/parameter.h>
+
+#include "../operator_common.h"
+namespace mxnet {
+namespace op {
+
+namespace beamsearch_set_finished {
+enum BeamsearchSetFinishedInputs {kDist, kScores, kFin, kOverMax};
+enum BeamsearchSetFinishedOutputs {kOut};
+}
+
+
+//template<int score_idx, int eos_idx>
+struct beamsearch_set_finished_forward {
+ template<typename DType, typename IType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType* in_data,
+ const DType* scores, const IType* fin, const IType* over_max,
+ const DType mask_val, const int score_idx, const int eos_idx,
+ int V) {
+ int j = i / V;
+ int k = i % V;
+ bool f = static_cast<bool>(fin[j]);
+ bool o = static_cast<bool>(over_max[j]);
+ bool s = k == score_idx;
+ bool e = k == eos_idx;
+ bool input = !f && (!o || e);
+ bool score = f && s;
+ //bool mask = !(input || score);
+ //out_data[i] = (input * in_data[i]) + (score * scores[j]) + (mask * mask_val);
+ if (input) out_data[i] = in_data[i];
+ else if (score) out_data[i] = scores[j];
+ else out_data[i] = mask_val;
+ }
+};
+
+struct BeamsearchSetFinishedParam : public dmlc::Parameter<BeamsearchSetFinishedParam> {
+ int score_idx;
+ int eos_idx;
+ float mask_val;
+ DMLC_DECLARE_PARAMETER(BeamsearchSetFinishedParam) {
+ DMLC_DECLARE_FIELD(score_idx)
+ .set_default(0)
+ .describe("Index to set the score of finished beams.");
+ DMLC_DECLARE_FIELD(eos_idx)
+ .describe("Index of the EOS token.");
+ DMLC_DECLARE_FIELD(mask_val)
+ .set_default(std::numeric_limits<float>::lowest())
+ .describe("Padding value used to mask out unwanted tokens in beams.");
+ }
+};
+
+inline bool BeamsearchSetFinishedShape(const nnvm::NodeAttrs& attrs,
+ mxnet::ShapeVector* in_attrs,
+ mxnet::ShapeVector* out_attrs) {
+ const BeamsearchSetFinishedParam& param = nnvm::get<BeamsearchSetFinishedParam>(attrs.parsed);
+ CHECK_EQ(in_attrs->size(), 4U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ auto dist = in_attrs->at(beamsearch_set_finished::kDist);
+ auto scores = in_attrs->at(beamsearch_set_finished::kScores);
+ auto fin = in_attrs->at(beamsearch_set_finished::kFin);
+ auto over_max = in_attrs->at(beamsearch_set_finished::kOverMax);
+ CHECK_EQ(dist.ndim(), 2U);
+ CHECK_EQ(scores.ndim(), 2U);
+ CHECK_EQ(fin.ndim(), 1U);
+ CHECK_EQ(over_max.ndim(), 1U);
+
+ CHECK_EQ(dist[0], scores[0]);
+ CHECK_EQ(dist[0], fin[0]);
+ CHECK_EQ(dist[0], over_max[0]);
+ CHECK_EQ(scores[1], 1);
+
+ mxnet::TShape score_shape(dist.ndim(), -1);
+ score_shape[0] = dist[0];
+ score_shape[1] = 1;
+
+ mxnet::TShape bool_shape(dist.ndim() - 1, -1);
+ bool_shape[0] = dist[0];
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, dist);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(beamsearch_set_finished::kOut));
+ SHAPE_ASSIGN_CHECK(*in_attrs, 1, score_shape);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 2, bool_shape);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 3, bool_shape);
+
+ return true;
+}
+
+inline bool BeamsearchSetFinishedType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 4U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 1, (*out_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 2, mshadow::kInt32);
+ TYPE_ASSIGN_CHECK(*in_attrs, 3, mshadow::kInt32);
+ return (*in_attrs)[0] != -1 && (*in_attrs)[1] != -1;
+}
+
+template<typename xpu>
+void NoopGrad(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ LOG(FATAL) << "This operator should only be used for inference";
+}
+
+template<typename xpu>
+void BeamsearchSetFinishedForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ if (req[beamsearch_set_finished::kOut] == mxnet::kNullOp) return;
+ const BeamsearchSetFinishedParam& param = nnvm::get<BeamsearchSetFinishedParam>(attrs.parsed);
+ CHECK_EQ(inputs.size(), 4U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+
+ const mxnet::TShape& out_shape = outputs[beamsearch_set_finished::kOut].shape_;
+ const mxnet::TShape& batch_beam_shape = inputs[beamsearch_set_finished::kFin].shape_;
+
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ using namespace mxnet_op;
+ MSHADOW_TYPE_SWITCH(outputs[beamsearch_set_finished::kOut].type_flag_, DType, {
+ MSHADOW_TYPE_SWITCH(inputs[beamsearch_set_finished::kFin].type_flag_, IType, {
Review comment:
seems like you are limiting the data type of `Fin` and `OverMax` to `int32` only in the `InferType` function, you then do not need a type switch here for those guys, you can simply do:
```c++
MSHADOW_TYPE_SWITCH(outputs[beamsearch_set_finished::kOut].type_flag_, DType, {
DType mask_val = param.mask_val;
const int score_idx = param.score_idx;
const int eos_idx = param.eos_idx;
Kernel<beamsearch_set_finished_forward, xpu>::Launch(
s, out_shape.Size(),
outputs[beamsearch_set_finished::kOut].dptr<DType>(),
inputs[beamsearch_set_finished::kDist].dptr<DType>(),
inputs[beamsearch_set_finished::kScores].dptr<DType>(),
inputs[beamsearch_set_finished::kFin].dptr<int>(),
inputs[beamsearch_set_finished::kOverMax].dptr<int>(),
mask_val, score_idx, eos_idx, out_shape.Size()/batch_beam_shape.Size());
});
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] haojin2 commented on a change in pull request
#17311: Added beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
haojin2 commented on a change in pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#discussion_r367028141
##########
File path: src/operator/contrib/beamsearch_set_finished.cc
##########
@@ -0,0 +1,66 @@
+#include <mxnet/base.h>
+#include "./beamsearch_set_finished-inl.h"
+#include "../tensor/elemwise_unary_op.h"
+
+namespace mxnet {
+namespace op {
+
+DMLC_REGISTER_PARAMETER(BeamsearchSetFinishedParam);
+
+NNVM_REGISTER_OP(_contrib_beamsearch_set_finished)
+.describe(R"code(Sets finished beams of the beam to a mask value (aside from the score index) and forces beams at max length to output the EOS sequence.
+
+Returns an array of the same shape of the input data array and the same values except for the designated rows whose elements are to be masked or be replaced by beam scores/EOS probabilities.
+
+Example::
+
+ x = [[ -1., -2., -3., -4.],
+ [ -5., -6., -7., -8.],
+ [ -9., -10., -11., -12.],
+ [-13., -14., -15., -16.]]
+
+ scores = [[-17.],
+ [-18.],
+ [-19.],
+ [-20.]]
+
+ finished = [0, 1, 0, 1]
+
+ over_max = [0, 0, 1, 1]
+
+ beamsearch_set_finished(x, scores, finished, over_max, score_idx=0,
+ eos_idx=2, mask_val=-1e15) = [[ -1., -2., -3., -4.],
+ [ -18., -1e15, -1e15, -1e15],
+ [-1e15, -1e15, -11., -1e15],
+ [ -20., -1e15, -1e15, -1e15]]
+
+.. Note::
+ This operator only supports forward propagation. DO NOT use it in training.
+
+)code")
+.set_attr_parser(ParamParser<BeamsearchSetFinishedParam>)
+.set_num_inputs(4)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data", "scores", "finished", "over_max"}; })
+.set_attr<nnvm::FListOutputNames>("FListOutputNames", [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"output"}; })
+.set_attr<FInferShape>("FInferShape", BeamsearchSetFinishedShape)
+.set_attr<nnvm::FInferType>("FInferType", BeamsearchSetFinishedType)
+.set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs) {
+ return std::vector<std::pair<int, int>>{{0, 0}}; })
+.set_attr<FCompute>("FCompute<cpu>", BeamsearchSetFinishedForward<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"beamsearch_noop_grad"})
+.add_argument("data", "NDArray-or-Symbol", "Input distribution of tokens")
+.add_argument("scores", "NDArray-or-Symbol", "Running scores for the sequences")
+.add_argument("finished", "NDArray-or-Symbol", "Finished beams")
+.add_argument("over_max", "NDArray-or-Symbol", "Beams at or exceeding maximum length")
+.add_arguments(BeamsearchSetFinishedParam::__FIELDS__());
+
+NNVM_REGISTER_OP(_contrib_beamsearch_noop_grad)
+.set_num_inputs(1)
+.set_num_outputs(4)
+.set_attr<FCompute>("FCompute<cpu>", NoopGrad<cpu>)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true);
+}
Review comment:
one more blank line below.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] blchu commented on issue #17311: Added
beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
blchu commented on issue #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#issuecomment-580003435
Currently there isn't support for fusing concatenation and broadcasting operators, though I'm not sure about the feasibility of doing so in the future.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] blchu closed pull request #17311: Added
beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
blchu closed pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] blchu edited a comment on issue #17311: Added
beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
blchu edited a comment on issue #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#issuecomment-580003435
Currently there isn't support for fusing concatenation and broadcasting operators, though I'm not sure about the feasibility of doing so in the future. From what I understand doing so is not trivial.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] blchu commented on a change in pull request
#17311: Added beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
blchu commented on a change in pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#discussion_r367184172
##########
File path: src/operator/contrib/beamsearch_set_finished-inl.h
##########
@@ -0,0 +1,148 @@
+#include <dmlc/parameter.h>
+
+#include "../operator_common.h"
+namespace mxnet {
+namespace op {
+
+namespace beamsearch_set_finished {
+enum BeamsearchSetFinishedInputs {kDist, kScores, kFin, kOverMax};
+enum BeamsearchSetFinishedOutputs {kOut};
+}
+
+
+//template<int score_idx, int eos_idx>
+struct beamsearch_set_finished_forward {
+ template<typename DType, typename IType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType* in_data,
+ const DType* scores, const IType* fin, const IType* over_max,
+ const DType mask_val, const int score_idx, const int eos_idx,
+ int V) {
+ int j = i / V;
+ int k = i % V;
+ bool f = static_cast<bool>(fin[j]);
+ bool o = static_cast<bool>(over_max[j]);
+ bool s = k == score_idx;
+ bool e = k == eos_idx;
+ bool input = !f && (!o || e);
+ bool score = f && s;
+ //bool mask = !(input || score);
+ //out_data[i] = (input * in_data[i]) + (score * scores[j]) + (mask * mask_val);
+ if (input) out_data[i] = in_data[i];
+ else if (score) out_data[i] = scores[j];
+ else out_data[i] = mask_val;
+ }
+};
+
+struct BeamsearchSetFinishedParam : public dmlc::Parameter<BeamsearchSetFinishedParam> {
+ int score_idx;
+ int eos_idx;
+ float mask_val;
+ DMLC_DECLARE_PARAMETER(BeamsearchSetFinishedParam) {
+ DMLC_DECLARE_FIELD(score_idx)
+ .set_default(0)
+ .describe("Index to set the score of finished beams.");
+ DMLC_DECLARE_FIELD(eos_idx)
+ .describe("Index of the EOS token.");
+ DMLC_DECLARE_FIELD(mask_val)
+ .set_default(std::numeric_limits<float>::lowest())
+ .describe("Padding value used to mask out unwanted tokens in beams.");
+ }
+};
+
+inline bool BeamsearchSetFinishedShape(const nnvm::NodeAttrs& attrs,
+ mxnet::ShapeVector* in_attrs,
+ mxnet::ShapeVector* out_attrs) {
+ const BeamsearchSetFinishedParam& param = nnvm::get<BeamsearchSetFinishedParam>(attrs.parsed);
+ CHECK_EQ(in_attrs->size(), 4U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ auto dist = in_attrs->at(beamsearch_set_finished::kDist);
+ auto scores = in_attrs->at(beamsearch_set_finished::kScores);
+ auto fin = in_attrs->at(beamsearch_set_finished::kFin);
+ auto over_max = in_attrs->at(beamsearch_set_finished::kOverMax);
+ CHECK_EQ(dist.ndim(), 2U);
+ CHECK_EQ(scores.ndim(), 2U);
+ CHECK_EQ(fin.ndim(), 1U);
+ CHECK_EQ(over_max.ndim(), 1U);
+
+ CHECK_EQ(dist[0], scores[0]);
+ CHECK_EQ(dist[0], fin[0]);
+ CHECK_EQ(dist[0], over_max[0]);
+ CHECK_EQ(scores[1], 1);
+
+ mxnet::TShape score_shape(dist.ndim(), -1);
+ score_shape[0] = dist[0];
+ score_shape[1] = 1;
+
+ mxnet::TShape bool_shape(dist.ndim() - 1, -1);
+ bool_shape[0] = dist[0];
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, dist);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(beamsearch_set_finished::kOut));
+ SHAPE_ASSIGN_CHECK(*in_attrs, 1, score_shape);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 2, bool_shape);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 3, bool_shape);
+
+ return true;
+}
+
+inline bool BeamsearchSetFinishedType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 4U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 1, (*out_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 2, mshadow::kInt32);
+ TYPE_ASSIGN_CHECK(*in_attrs, 3, mshadow::kInt32);
+ return (*in_attrs)[0] != -1 && (*in_attrs)[1] != -1;
+}
+
+template<typename xpu>
+void NoopGrad(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ LOG(FATAL) << "This operator should only be used for inference";
+}
+
+template<typename xpu>
+void BeamsearchSetFinishedForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ if (req[beamsearch_set_finished::kOut] == mxnet::kNullOp) return;
+ const BeamsearchSetFinishedParam& param = nnvm::get<BeamsearchSetFinishedParam>(attrs.parsed);
+ CHECK_EQ(inputs.size(), 4U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+
+ const mxnet::TShape& out_shape = outputs[beamsearch_set_finished::kOut].shape_;
+ const mxnet::TShape& batch_beam_shape = inputs[beamsearch_set_finished::kFin].shape_;
+
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ using namespace mxnet_op;
+ MSHADOW_TYPE_SWITCH(outputs[beamsearch_set_finished::kOut].type_flag_, DType, {
+ MSHADOW_TYPE_SWITCH(inputs[beamsearch_set_finished::kFin].type_flag_, IType, {
Review comment:
I see, I'll change that then
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] sxjscience commented on issue #17311: Added
beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
sxjscience commented on issue #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#issuecomment-578335377
Is it possible to automatically fuse these operators via the recently introduced point-wise fusion?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-mxnet] haojin2 commented on a change in pull request
#17311: Added beamsearch_set_finished Operator
Posted by GitBox <gi...@apache.org>.
haojin2 commented on a change in pull request #17311: Added beamsearch_set_finished Operator
URL: https://github.com/apache/incubator-mxnet/pull/17311#discussion_r367027893
##########
File path: src/operator/contrib/beamsearch_set_finished-inl.h
##########
@@ -0,0 +1,148 @@
+#include <dmlc/parameter.h>
+
+#include "../operator_common.h"
+namespace mxnet {
+namespace op {
+
+namespace beamsearch_set_finished {
+enum BeamsearchSetFinishedInputs {kDist, kScores, kFin, kOverMax};
+enum BeamsearchSetFinishedOutputs {kOut};
+}
+
+
+//template<int score_idx, int eos_idx>
+struct beamsearch_set_finished_forward {
+ template<typename DType, typename IType>
+ MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType* in_data,
+ const DType* scores, const IType* fin, const IType* over_max,
+ const DType mask_val, const int score_idx, const int eos_idx,
+ int V) {
+ int j = i / V;
+ int k = i % V;
+ bool f = static_cast<bool>(fin[j]);
+ bool o = static_cast<bool>(over_max[j]);
+ bool s = k == score_idx;
+ bool e = k == eos_idx;
+ bool input = !f && (!o || e);
+ bool score = f && s;
+ //bool mask = !(input || score);
+ //out_data[i] = (input * in_data[i]) + (score * scores[j]) + (mask * mask_val);
+ if (input) out_data[i] = in_data[i];
+ else if (score) out_data[i] = scores[j];
+ else out_data[i] = mask_val;
+ }
+};
+
+struct BeamsearchSetFinishedParam : public dmlc::Parameter<BeamsearchSetFinishedParam> {
+ int score_idx;
+ int eos_idx;
+ float mask_val;
+ DMLC_DECLARE_PARAMETER(BeamsearchSetFinishedParam) {
+ DMLC_DECLARE_FIELD(score_idx)
+ .set_default(0)
+ .describe("Index to set the score of finished beams.");
+ DMLC_DECLARE_FIELD(eos_idx)
+ .describe("Index of the EOS token.");
+ DMLC_DECLARE_FIELD(mask_val)
+ .set_default(std::numeric_limits<float>::lowest())
+ .describe("Padding value used to mask out unwanted tokens in beams.");
+ }
+};
+
+inline bool BeamsearchSetFinishedShape(const nnvm::NodeAttrs& attrs,
+ mxnet::ShapeVector* in_attrs,
+ mxnet::ShapeVector* out_attrs) {
+ const BeamsearchSetFinishedParam& param = nnvm::get<BeamsearchSetFinishedParam>(attrs.parsed);
+ CHECK_EQ(in_attrs->size(), 4U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ auto dist = in_attrs->at(beamsearch_set_finished::kDist);
+ auto scores = in_attrs->at(beamsearch_set_finished::kScores);
+ auto fin = in_attrs->at(beamsearch_set_finished::kFin);
+ auto over_max = in_attrs->at(beamsearch_set_finished::kOverMax);
+ CHECK_EQ(dist.ndim(), 2U);
+ CHECK_EQ(scores.ndim(), 2U);
+ CHECK_EQ(fin.ndim(), 1U);
+ CHECK_EQ(over_max.ndim(), 1U);
+
+ CHECK_EQ(dist[0], scores[0]);
+ CHECK_EQ(dist[0], fin[0]);
+ CHECK_EQ(dist[0], over_max[0]);
+ CHECK_EQ(scores[1], 1);
+
+ mxnet::TShape score_shape(dist.ndim(), -1);
+ score_shape[0] = dist[0];
+ score_shape[1] = 1;
+
+ mxnet::TShape bool_shape(dist.ndim() - 1, -1);
+ bool_shape[0] = dist[0];
+
+ SHAPE_ASSIGN_CHECK(*out_attrs, 0, dist);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(beamsearch_set_finished::kOut));
+ SHAPE_ASSIGN_CHECK(*in_attrs, 1, score_shape);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 2, bool_shape);
+ SHAPE_ASSIGN_CHECK(*in_attrs, 3, bool_shape);
+
+ return true;
+}
+
+inline bool BeamsearchSetFinishedType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 4U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 1, (*out_attrs)[0]);
+ TYPE_ASSIGN_CHECK(*in_attrs, 2, mshadow::kInt32);
+ TYPE_ASSIGN_CHECK(*in_attrs, 3, mshadow::kInt32);
+ return (*in_attrs)[0] != -1 && (*in_attrs)[1] != -1;
+}
+
+template<typename xpu>
+void NoopGrad(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ LOG(FATAL) << "This operator should only be used for inference";
+}
+
+template<typename xpu>
+void BeamsearchSetFinishedForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ if (req[beamsearch_set_finished::kOut] == mxnet::kNullOp) return;
+ const BeamsearchSetFinishedParam& param = nnvm::get<BeamsearchSetFinishedParam>(attrs.parsed);
+ CHECK_EQ(inputs.size(), 4U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+
+ const mxnet::TShape& out_shape = outputs[beamsearch_set_finished::kOut].shape_;
+ const mxnet::TShape& batch_beam_shape = inputs[beamsearch_set_finished::kFin].shape_;
+
+ mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
+ using namespace mxnet_op;
+ MSHADOW_TYPE_SWITCH(outputs[beamsearch_set_finished::kOut].type_flag_, DType, {
+ MSHADOW_TYPE_SWITCH(inputs[beamsearch_set_finished::kFin].type_flag_, IType, {
Review comment:
seems like you are limiting the data type of `Fin` and `OverMax` to `int32` only in the `InferType` function, you then do not need a type switch here for those guys, you can simply do:
```c++
MSHADOW_TYPE_SWITCH(outputs[beamsearch_set_finished::kOut].type_flag_, DType, {
DType mask_val = param.mask_val;
const int score_idx = param.score_idx;
const int eos_idx = param.eos_idx;
Kernel<beamsearch_set_finished_forward, xpu>::Launch(
s, out_shape.Size(),
outputs[beamsearch_set_finished::kOut].dptr<DType>(),
inputs[beamsearch_set_finished::kDist].dptr<DType>(),
inputs[beamsearch_set_finished::kScores].dptr<DType>(),
inputs[beamsearch_set_finished::kFin].dptr<int32>(),
inputs[beamsearch_set_finished::kOverMax].dptr<int32>(),
mask_val, score_idx, eos_idx, out_shape.Size()/batch_beam_shape.Size());
});
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services