[TOPI][Relay][TensorFlow] Add OneHot operator #3781
Conversation
soiferj
changed the title
python/tvm/relay/op/transform.py
Outdated
|
|
||
| indices = [1., 2., 3.] | ||
|
|
||
| relay.one_hot(indices, 2) = |
There was a problem hiding this comment.
Should the depth be 3 in this example?
There was a problem hiding this comment.
Yes, you're right, I'll fix that.
python/tvm/relay/op/transform.py
Outdated
| -------- | ||
| .. code-block:: python | ||
|
|
||
| indices = [1., 2., 3.] |
There was a problem hiding this comment.
Should indices be [0, 1, 2] in this example?
topi/include/topi/transform.h
Outdated
| } | ||
|
|
||
| auto idx = iter_vars[iter_vars.size() - 1]; | ||
| return tvm::if_then_else(indices(outer_indices) == idx, 1, 0); |
There was a problem hiding this comment.
You should use ir::Select::make here for better performance
There was a problem hiding this comment.
Do you require indices to be int type? Otherwise, this line could fail if indices.dtype is float32 without type cast.
You should add a test case for indices with float32 type.
There was a problem hiding this comment.
Can you expand on how to use ir::Select::make here? And yes, indices must be an int type.
There was a problem hiding this comment.
Just replace tvm::if_then_else by ir::Select::make. You can find use case at https://github.com/dmlc/tvm/blob/master/topi/include/topi/nn/pooling.h#L265.
|
Could you also add |
|
@icemelon9 sure, I can do that. The tricky part about on_value and off_value is that it can be either a float or an int, meaning that OneHotAttrs will not necessarily have the correct data type. What do you think is the best way to handle this? Should I have the on_value and off_value be of type float in OneHotAttrs, and do any casting in the TOPI definition? |
|
You can just cast it to the dtype in the attr, and put on_value and off_value as Expr in argument not in the attribute. You can check out implementation of full op: |
|
Added support for on_value, off_value, axis, and dtype. Let me know what you think! |
|
@icemelon9 would you mind looking at the updates? |
|
lgtm |
|
Added to the mxnet frontend |
|
It seems not necessary to pass on_value and off_value as tvm.relay.Expr, why not just treat them as attributes? That way, we can avoid allocating two more tensors. |
|
The on_value and off_value can be any data type. Treating them as tvm.relay.Expr makes the casting to the correct data type easier. |
* Add one-hot to Relay * topi implementation * Working * add topi test * Add TF test * Fix check * fix linting issues * fix documentation * Fix documentation * Add support for on_value, off_value, axis, dtype * Add full support for axis * Fix compute and update test_forward * Move on_value and off_value to inputs * Add topi test * Update tests * Update docs * Fix style * re-enable tests * Add one_hot to mxnet converter
* Add one-hot to Relay * topi implementation * Working * add topi test * Add TF test * Fix check * fix linting issues * fix documentation * Fix documentation * Add support for on_value, off_value, axis, dtype * Add full support for axis * Fix compute and update test_forward * Move on_value and off_value to inputs * Add topi test * Update tests * Update docs * Fix style * re-enable tests * Add one_hot to mxnet converter
* Add one-hot to Relay * topi implementation * Working * add topi test * Add TF test * Fix check * fix linting issues * fix documentation * Fix documentation * Add support for on_value, off_value, axis, dtype * Add full support for axis * Fix compute and update test_forward * Move on_value and off_value to inputs * Add topi test * Update tests * Update docs * Fix style * re-enable tests * Add one_hot to mxnet converter
Add one-hot operator where on_value is 1 and off_value is 0. This implementation is enough to support most scenarios. We can add support for on_value and off_value at a later time if it is needed.
@jroesch @icemelon9 would you be able to take a look?
adding @Laurawly and @Huyuwei