| @@ -53,8 +53,13 @@ Status InsertReshapeIfNeed(const NodePtr &node) { | |||||
| node->GetName().c_str(), src_anchor->GetIdx(), dst_node->GetName().c_str(), dst_anchor->GetIdx()); | node->GetName().c_str(), src_anchor->GetIdx(), dst_node->GetName().c_str(), dst_anchor->GetIdx()); | ||||
| GE_CHECK_NOTNULL(dst_node); | GE_CHECK_NOTNULL(dst_node); | ||||
| GE_CHECK_NOTNULL(dst_node->GetOpDesc()); | GE_CHECK_NOTNULL(dst_node->GetOpDesc()); | ||||
| auto dst_tensor = dst_node->GetOpDesc()->GetInputDescPtr(dst_anchor->GetIdx()); | |||||
| auto dst_tensor = dst_node->GetOpDesc()->MutableInputDesc(dst_anchor->GetIdx()); | |||||
| GE_CHECK_NOTNULL(dst_tensor); | GE_CHECK_NOTNULL(dst_tensor); | ||||
| if (dst_node->GetType() == NETOUTPUT) { | |||||
| // NetOutput shape must be continuous, do not insert reshape node before NetOutput. | |||||
| dst_tensor->SetShape(src_tensor->GetShape()); | |||||
| continue; | |||||
| } | |||||
| bool is_need_insert_reshape = src_tensor->GetShape().GetDims() != UNKNOWN_RANK && | bool is_need_insert_reshape = src_tensor->GetShape().GetDims() != UNKNOWN_RANK && | ||||
| dst_tensor->GetShape().GetDims() != UNKNOWN_RANK && | dst_tensor->GetShape().GetDims() != UNKNOWN_RANK && | ||||
| src_tensor->GetShape().GetDims() != dst_tensor->GetShape().GetDims(); | src_tensor->GetShape().GetDims() != dst_tensor->GetShape().GetDims(); | ||||
| @@ -680,6 +680,7 @@ set(PASS_TEST_FILES | |||||
| "graph/passes/assert_pass_unittest.cc" | "graph/passes/assert_pass_unittest.cc" | ||||
| "graph/passes/dropout_pass_unittest.cc" | "graph/passes/dropout_pass_unittest.cc" | ||||
| "graph/passes/unused_const_pass_unittest.cc" | "graph/passes/unused_const_pass_unittest.cc" | ||||
| "graph/passes/reshape_recovery_pass_unittest.cc" | |||||
| "graph/passes/reshape_remove_pass_unittest.cc" | "graph/passes/reshape_remove_pass_unittest.cc" | ||||
| "graph/passes/resource_pair_control_pass_unittest.cc" | "graph/passes/resource_pair_control_pass_unittest.cc" | ||||
| "graph/passes/trans_op_breadth_fusion_pass_unittest.cc" | "graph/passes/trans_op_breadth_fusion_pass_unittest.cc" | ||||
| @@ -0,0 +1,55 @@ | |||||
| /** | |||||
| * Copyright 2019-2020 Huawei Technologies Co., Ltd | |||||
| * | |||||
| * Licensed under the Apache License, Version 2.0 (the "License"); | |||||
| * you may not use this file except in compliance with the License. | |||||
| * You may obtain a copy of the License at | |||||
| * | |||||
| * http://www.apache.org/licenses/LICENSE-2.0 | |||||
| * | |||||
| * Unless required by applicable law or agreed to in writing, software | |||||
| * distributed under the License is distributed on an "AS IS" BASIS, | |||||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |||||
| * See the License for the specific language governing permissions and | |||||
| * limitations under the License. | |||||
| */ | |||||
| #include "graph/passes/reshape_recovery_pass.h" | |||||
| #include <gtest/gtest.h> | |||||
| #include <set> | |||||
| #include <string> | |||||
| #include "graph_builder_utils.h" | |||||
| namespace ge { | |||||
| class UtestReshapeRecoveryPass : public testing::Test { | |||||
| protected: | |||||
| void SetUp() {} | |||||
| void TearDown() {} | |||||
| }; | |||||
| namespace { | |||||
| /// data | |||||
| /// | | |||||
| /// transdata | |||||
| /// | | |||||
| /// netoutput | |||||
| ComputeGraphPtr BuilderGraph1() { | |||||
| ut::GraphBuilder builder = ut::GraphBuilder("g1"); | |||||
| auto data = builder.AddNode("data", "Data", 0, 1, FORMAT_ND, DT_FLOAT, {1, 3, 16, 16}); | |||||
| auto transdata = builder.AddNode("transdata", "Transdata", 1, 1, FORMAT_ND, DT_FLOAT, {1, 3, 16, 16}); | |||||
| auto netoutput = builder.AddNode("netoutput", "Netoutput", 1, 0, FORMAT_ND, DT_FLOAT, {1, 3, 16, 16}); | |||||
| builder.AddDataEdge(data, 0, transdata, 0); | |||||
| builder.AddDataEdge(transdata, 0, netoutput, 0); | |||||
| return builder.GetGraph(); | |||||
| } | |||||
| } | |||||
| TEST_F(UtestReshapeRecoveryPass, test_reshape_recovery) { | |||||
| auto graph = BuilderGraph1(); | |||||
| ReshapeRecoveryPass pass; | |||||
| EXPECT_EQ(pass.Run(graph), SUCCESS); | |||||
| } | |||||
| } | |||||