yankf21
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add NHWC format requirement of topk ops

This commit is contained in:
liuwenhao4 2020-12-09 17:53:39 +08:00
parent a9fcea2ca6
commit 1957afa243
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mindspore/lite/src/ops/topk.cc
View File

@ -60,6 +60,10 @@ int TopK::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> output
}
auto input = inputs_.front();
MS_ASSERT(input != nullptr);
if (input->format() != schema::Format::Format_NHWC) {
MS_LOG(ERROR) << "topk only support NHWC now!";
return RET_FORMAT_ERR;
}
auto output0 = outputs_.front();
MS_ASSERT(output0 != nullptr);
auto output1 = outputs_.at(1);

3
mindspore/lite/tools/common/node_util.cc
View File

@ -50,7 +50,8 @@ static const std::vector<schema::PrimitiveType> nhwcOpList = {
schema::PrimitiveType_BiasAdd,
schema::PrimitiveType_InstanceNorm,
schema::PrimitiveType_SpaceToDepth,
schema::PrimitiveType_DepthToSpace};
schema::PrimitiveType_DepthToSpace,
schema::PrimitiveType_TopK};
static const std::vector<schema::PrimitiveType> nhwcOpAllInputList = {
#ifdef SUPPORT_TRAIN