[NNVM][TOPI] Add gradients for broadcast_* ops #1234
Conversation
|
@kevinthesun @nishi-t can you help review this? |
nnvm/src/top/tensor/reduce.cc
Outdated
| } | ||
| } else { | ||
| --j; | ||
| } |
There was a problem hiding this comment.
Is it possible to make L271-281 more simple ?
For example:
if ( j < 0 ) {
r_axes.push_back(i);
squeeze_axes.push_back(i);
} else if ( ishpae[i] != oshape[j]) {
r_axes.push_back(i);
if (oshape[j] == 1) {
keepdims = true;
--j;
}
} else {
--j;
}There was a problem hiding this comment.
Just a bit. I changed the index variable names and added some comments to help make it clearer.
nnvm/src/top/tensor/reduce.cc
Outdated
| NNVM_REGISTER_COLLAPSE_OP(sum) | ||
| .describe(R"code(Reduces lhs to the shape of rhs via sum)code" NNVM_ADD_FILELINE) | ||
| .set_attr<FTVMCompute>( | ||
| "FTVMCompute", [](const NodeAttrs& attrs, |
There was a problem hiding this comment.
Can we move computation part into tvm? It's valuable to pack it first in tvm and then directly call in nnvm.
There was a problem hiding this comment.
Apparently not. Moving the entire op into topi yields RuntimeError: Not yet support ewise after reduce. I guess I could turn the entire thing into a single topi op, but that defeats the purpose of writing composable functions.
nnvm/src/top/tensor/reduce.cc
Outdated
| .add_arguments(ReduceParam::__FIELDS__()) \ | ||
| .set_attr_parser(AxesParamParser<ReduceParam>) \ | ||
| .set_attr<FInferShape>("FInferShape", CollapseShape) \ | ||
| .set_attr<FInferType>("FInferType", ElemwiseType<2, 1>) \ |
There was a problem hiding this comment.
We need FCorrectLayout attribute to for correct layout pass. An example for elemwise binary op https://github.com/dmlc/tvm/blob/master/nnvm/src/top/elemwise_op_common.h#L301.
| @@ -130,7 +134,61 @@ def test_multi_loss_graph_gradients(): | |||
| # test reverse infer type for label | |||
| assert grad_g.apply('InferType').json_attr('dtype_num_unknown_nodes') == 0 | |||
|
|
|||
| def _run_op_grad(op): | |||
There was a problem hiding this comment.
For operator unit test, we can add them under compiler/test_top_levelx.py, such as https://github.com/dmlc/tvm/blob/master/nnvm/tests/python/compiler/test_top_level1.py. Helper function supporting backward is implemented so that gradient can be tested here.
tqchen
added
status: review in progress
status: need update
|
@nhynes can you act on the review comments? |
|
@tqchen @kevinthesun @nishi-t How can I get this PR to work with cuda? The issue is supporting no-op collapse. How should I proceed? This PR fixes several bugs with reduce, so it'd be a shame to see it go to waste. |
|
The easiest way might be to patch schedule_reduce to support no-reduce-axis case, and call schedule_injective in that case, which is always safe |
tqchen
added
status: accepted
and removed
status: need update
This PR
broadcast_*ops by adding acollapse_sumopMakeReduceTargetShapewhich prevents reduction to scalar