1 //
2 // CPUOneHot.cpp
3 // MNN
4 //
5 // Created by MNN on 2019/11/29.
6 // Copyright © 2018, Alibaba Group Holding Limited
7 //
8
9 #include "backend/cpu/CPUOneHot.hpp"
10 #include "backend/cpu/CPUBackend.hpp"
11
12 namespace MNN {
13
14 template <typename T>
OneHotImpl(int depth,int outerSize,int innerSize,const int * indices,const Tensor * onValueTensor,const Tensor * offValueTensor,Tensor * outputTensor)15 void OneHotImpl(int depth, int outerSize, int innerSize, const int* indices, const Tensor* onValueTensor,
16 const Tensor* offValueTensor, Tensor* outputTensor) {
17 const T onValue = onValueTensor->host<T>()[0];
18 const T offValue = offValueTensor->host<T>()[0];
19 T* outputPtr = outputTensor->host<T>();
20
21 for (int i = 0; i < outerSize; ++i) {
22 for (int j = 0; j < depth; ++j) {
23 for (int k = 0; k < innerSize; ++k) {
24 auto index = indices[i * innerSize + k];
25 if (index == j) {
26 *outputPtr = onValue;
27 } else {
28 *outputPtr = offValue;
29 }
30 outputPtr++;
31 }
32 }
33 }
34 }
35
onExecute(const std::vector<Tensor * > & inputs,const std::vector<Tensor * > & outputs)36 ErrorCode CPUOneHot::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
37 auto indices = inputs[0];
38 auto depthTensor = inputs[1];
39 auto onValueTensor = inputs[2];
40 auto offValueTensor = inputs[3];
41
42 int axis = mAxis;
43 if (axis < 0) {
44 axis += outputs[0]->dimensions();
45 }
46 int outerSize = 1;
47 for (int i = 0; i < axis; ++i) {
48 outerSize *= indices->length(i);
49 }
50 const int depth = depthTensor->host<int>()[0];
51 const int innerSize = indices->elementSize() / outerSize;
52 const auto indicesPtr = indices->host<int>();
53
54 auto dataType = onValueTensor->getType();
55 MNN_ASSERT(offValueTensor->getType() == dataType);
56
57 if (dataType == halide_type_of<float>()) {
58 OneHotImpl<float>(depth, outerSize, innerSize, indicesPtr, onValueTensor, offValueTensor, outputs[0]);
59 } else if (dataType == halide_type_of<int>()) {
60 OneHotImpl<int>(depth, outerSize, innerSize, indicesPtr, onValueTensor, offValueTensor, outputs[0]);
61 } else {
62 return NOT_SUPPORT;
63 }
64 return NO_ERROR;
65 }
66
67 class CPUOneHotCreator : public CPUBackend::Creator {
68 public:
onCreate(const std::vector<Tensor * > & inputs,const std::vector<Tensor * > & outputs,const MNN::Op * op,Backend * backend) const69 virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
70 const MNN::Op* op, Backend* backend) const override {
71 return new CPUOneHot(backend, op->main_as_OneHotParam()->axis());
72 }
73 };
74
75 REGISTER_CPU_OP_CREATOR(CPUOneHotCreator, OpType_OneHot);
76
77 } // namespace MNN
78