mirror of https://github.com/alibaba/MNN.git
78 lines
2.5 KiB
C++
78 lines
2.5 KiB
C++
//
|
|
// CPUOneHot.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/11/29.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "backend/cpu/CPUOneHot.hpp"
|
|
#include "backend/cpu/CPUBackend.hpp"
|
|
|
|
namespace MNN {
|
|
|
|
template <typename T>
|
|
void OneHotImpl(int depth, int outerSize, int innerSize, const int* indices, const Tensor* onValueTensor,
|
|
const Tensor* offValueTensor, Tensor* outputTensor) {
|
|
const T onValue = onValueTensor->host<T>()[0];
|
|
const T offValue = offValueTensor->host<T>()[0];
|
|
T* outputPtr = outputTensor->host<T>();
|
|
|
|
for (int i = 0; i < outerSize; ++i) {
|
|
for (int j = 0; j < depth; ++j) {
|
|
for (int k = 0; k < innerSize; ++k) {
|
|
auto index = indices[i * innerSize + k];
|
|
if (index == j) {
|
|
*outputPtr = onValue;
|
|
} else {
|
|
*outputPtr = offValue;
|
|
}
|
|
outputPtr++;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
ErrorCode CPUOneHot::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto indices = inputs[0];
|
|
auto depthTensor = inputs[1];
|
|
auto onValueTensor = inputs[2];
|
|
auto offValueTensor = inputs[3];
|
|
|
|
int axis = mAxis;
|
|
if (axis < 0) {
|
|
axis += outputs[0]->dimensions();
|
|
}
|
|
int outerSize = 1;
|
|
for (int i = 0; i < axis; ++i) {
|
|
outerSize *= indices->length(i);
|
|
}
|
|
const int depth = depthTensor->host<int>()[0];
|
|
const int innerSize = indices->elementSize() / outerSize;
|
|
const auto indicesPtr = indices->host<int>();
|
|
|
|
auto dataType = onValueTensor->getType();
|
|
MNN_ASSERT(offValueTensor->getType() == dataType);
|
|
|
|
if (dataType == halide_type_of<float>()) {
|
|
OneHotImpl<float>(depth, outerSize, innerSize, indicesPtr, onValueTensor, offValueTensor, outputs[0]);
|
|
} else if (dataType == halide_type_of<int>()) {
|
|
OneHotImpl<int>(depth, outerSize, innerSize, indicesPtr, onValueTensor, offValueTensor, outputs[0]);
|
|
} else {
|
|
return NOT_SUPPORT;
|
|
}
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class CPUOneHotCreator : public CPUBackend::Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const MNN::Op* op, Backend* backend) const override {
|
|
return new CPUOneHot(backend, op->main_as_OneHotParam()->axis());
|
|
}
|
|
};
|
|
|
|
REGISTER_CPU_OP_CREATOR(CPUOneHotCreator, OpType_OneHot);
|
|
|
|
} // namespace MNN
|