提交 39ac3273 编写于 作者: K kai00

fix anf exporter

上级 02c92155
......@@ -241,6 +241,9 @@ void AnfExporter::SetOpInputNode(const CNodePtr &cnode, schema::MetaGraphT *meta
nodeIdMap[valueNode->fullname_with_scope()] = meta_graph->allTensors.size();
fbNode->inputIndex.emplace_back(meta_graph->allTensors.size());
meta_graph->allTensors.emplace_back(std::move(paramTensor));
} else if (value->isa<mindspore::ValueSequeue>()) {
MS_LOG(INFO) << "Value type is ValueSequence.";
break;
} else {
MS_LOG(ERROR) << "Not support value type , need add support.";
}
......
......@@ -21,6 +21,10 @@
#include "ir/primitive.h"
namespace mindspore::lite {
namespace {
constexpr int kReduceInputNum = 3;
constexpr int kReduceInputIndex = 2;
}
int mindspore::lite::AnfReduceMeanPopulater::Parse(CNodePtr cnodePtr, schema::CNodeT *node,
std::vector<schema::TensorT *> *outputs) {
auto p = GetCNodePrimitive(cnodePtr);
......@@ -28,7 +32,25 @@ int mindspore::lite::AnfReduceMeanPopulater::Parse(CNodePtr cnodePtr, schema::CN
attr->mode = schema::ReduceMode_ReduceMean;
attr->keepDims = GetValue<bool>(p->GetAttr("keep_dims"));
// attr->axes = GetValue<std::vector<int>>(p->GetAttr("shape"));
if (cnodePtr->inputs().size() == kReduceInputNum) {
auto inputNode = cnodePtr->input(kReduceInputIndex);
MS_ASSERT(inputNode != nullptr);
if (inputNode->isa<ValueNode>()) {
auto valueNode = inputNode->cast<ValueNodePtr>();
MS_ASSERT(valueNode != nullptr);
auto value = valueNode->value();
MS_ASSERT(value != nullptr);
if (value->isa<ValueTuple>()) {
auto valTuplPtr = dyn_cast<ValueTuple>(value);
MS_ASSERT(valTuplPtr != nullptr);
for (size_t i = 0; i < valTuplPtr->size(); i++) {
auto elem = dyn_cast<Int32Imm>((*valTuplPtr)[i]);
MS_ASSERT(elem != nullptr);
attr->axes.emplace_back(elem->value());
}
}
}
}
node->nodeType = schema::NodeType_CNode;
node->primitive = std::make_unique<schema::PrimitiveT>();
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册