yujianfeng
5 years ago
parent
11732f0ea2
commit
188d74f15e
@ -0,0 +1,83 @@
|
||||
/**
|
||||
* Copyright 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 "backend/optimizer/ascend/format_type/remove_internal_output.h"
|
||||
#include <memory>
|
||||
#include "backend/session/anf_runtime_algorithm.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace opt {
|
||||
namespace {
|
||||
bool UsedForOutputOnly(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
|
||||
MS_EXCEPTION_IF_NULL(func_graph);
|
||||
auto manager = func_graph->manager();
|
||||
MS_EXCEPTION_IF_NULL(manager);
|
||||
auto &node_users = manager->node_users();
|
||||
auto iter = node_users.find(node);
|
||||
if (iter == node_users.end()) {
|
||||
return false;
|
||||
}
|
||||
const auto &node_set = iter->second;
|
||||
for (const auto &node_index : node_set) {
|
||||
if (!AnfAlgo::CheckPrimitiveType(node_index.first, prim::kPrimMakeTuple)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
} // namespace
|
||||
const BaseRef RemoveInternalOutputTransOp::DefinePattern() const {
|
||||
VarPtr X = std::make_shared<Var>();
|
||||
auto prim = std::make_shared<Primitive>(kTransDataOpName);
|
||||
return VectorRef({prim, X});
|
||||
}
|
||||
|
||||
const BaseRef RemoveInternalOutputCast::DefinePattern() const {
|
||||
VarPtr X = std::make_shared<Var>();
|
||||
return VectorRef({prim::kPrimCast, X});
|
||||
}
|
||||
|
||||
const AnfNodePtr RemoveInternalOutput::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
|
||||
const EquivPtr &) const {
|
||||
MS_EXCEPTION_IF_NULL(func_graph);
|
||||
MS_EXCEPTION_IF_NULL(node);
|
||||
auto kernel_graph = func_graph->cast<KernelGraphPtr>();
|
||||
if (kernel_graph == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
if (!kernel_graph->IsInternalOutput(node)) {
|
||||
return nullptr;
|
||||
}
|
||||
if (!UsedForOutputOnly(func_graph, node)) {
|
||||
return nullptr;
|
||||
}
|
||||
auto cnode = node->cast<CNodePtr>();
|
||||
MS_EXCEPTION_IF_NULL(cnode);
|
||||
CheckCNodeInputSize(cnode, kTransOpInputNum);
|
||||
auto input_node = cnode->input(1);
|
||||
if (!AnfAlgo::CheckPrimitiveType(input_node, prim::kPrimTupleGetItem)) {
|
||||
kernel_graph->ReplaceInternalOutput(node, input_node);
|
||||
} else {
|
||||
auto tuple_getitem = input_node->cast<CNodePtr>();
|
||||
MS_EXCEPTION_IF_NULL(tuple_getitem);
|
||||
int idx = AnfAlgo::GetTupleGetItemOutIndex(tuple_getitem);
|
||||
AnfNodePtr real_input_node = AnfAlgo::GetTupleGetItemRealInput(tuple_getitem);
|
||||
kernel_graph->ReplaceInternalOutput(node, real_input_node, 0, idx);
|
||||
}
|
||||
return input_node;
|
||||
}
|
||||
} // namespace opt
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,51 @@
|
||||
/**
|
||||
* Copyright 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.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_FORMAT_TYPE_REMOVE_INTERNAL_OUTPUT_H_
|
||||
#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_FORMAT_TYPE_REMOVE_INTERNAL_OUTPUT_H_
|
||||
|
||||
#include <string>
|
||||
#include "backend/optimizer/common/optimizer.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace opt {
|
||||
class RemoveInternalOutput : public PatternProcessPass {
|
||||
public:
|
||||
explicit RemoveInternalOutput(const std::string &name, bool multigraph = true)
|
||||
: PatternProcessPass(name, multigraph) {}
|
||||
~RemoveInternalOutput() override = default;
|
||||
const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
|
||||
};
|
||||
|
||||
class RemoveInternalOutputTransOp : public RemoveInternalOutput {
|
||||
public:
|
||||
explicit RemoveInternalOutputTransOp(bool multigraph = true)
|
||||
: RemoveInternalOutput("remove_internal_output_trans_op", multigraph) {}
|
||||
~RemoveInternalOutputTransOp() override = default;
|
||||
const BaseRef DefinePattern() const override;
|
||||
};
|
||||
|
||||
class RemoveInternalOutputCast : public RemoveInternalOutput {
|
||||
public:
|
||||
explicit RemoveInternalOutputCast(bool multigraph = true)
|
||||
: RemoveInternalOutput("remove_internal_output_cast", multigraph) {}
|
||||
~RemoveInternalOutputCast() override = default;
|
||||
const BaseRef DefinePattern() const override;
|
||||
};
|
||||
} // namespace opt
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_FORMAT_TYPE_REMOVE_INTERNAL_OUTPUT_H_
|
||||
@ -0,0 +1,89 @@
|
||||
# Copyright 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.
|
||||
# ============================================================================
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import mindspore.context as context
|
||||
import mindspore.nn as nn
|
||||
from mindspore import Tensor
|
||||
from mindspore.nn import TrainOneStepCell, WithLossCell
|
||||
from mindspore.nn.optim import Momentum
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
|
||||
|
||||
|
||||
class LeNet(nn.Cell):
|
||||
def __init__(self):
|
||||
super(LeNet, self).__init__()
|
||||
self.relu = P.ReLU()
|
||||
self.batch_size = 32
|
||||
|
||||
self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=0, has_bias=False, pad_mode='valid')
|
||||
self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0, has_bias=False, pad_mode='valid')
|
||||
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
self.reshape = P.Reshape()
|
||||
self.fc1 = nn.Dense(400, 120)
|
||||
self.fc1.matmul.add_prim_attr("primitive_target", "CPU")
|
||||
self.fc1.bias_add.add_prim_attr("primitive_target", "CPU")
|
||||
self.fc2 = nn.Dense(120, 84)
|
||||
self.fc2.matmul.add_prim_attr("primitive_target", "CPU")
|
||||
self.fc2.bias_add.add_prim_attr("primitive_target", "CPU")
|
||||
self.fc3 = nn.Dense(84, 10)
|
||||
self.fc3.matmul.add_prim_attr("primitive_target", "CPU")
|
||||
self.fc3.bias_add.add_prim_attr("primitive_target", "CPU")
|
||||
|
||||
def construct(self, input_x):
|
||||
output = self.conv1(input_x)
|
||||
output = self.relu(output)
|
||||
output = self.pool(output)
|
||||
output = self.conv2(output)
|
||||
output = self.relu(output)
|
||||
output = self.pool(output)
|
||||
output = self.reshape(output, (self.batch_size, -1))
|
||||
output = self.fc1(output)
|
||||
output = self.relu(output)
|
||||
output = self.fc2(output)
|
||||
output = self.relu(output)
|
||||
output = self.fc3(output)
|
||||
return output
|
||||
|
||||
|
||||
def train(net, data, label):
|
||||
learning_rate = 0.01
|
||||
momentum = 0.9
|
||||
|
||||
optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate, momentum)
|
||||
criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True)
|
||||
net_with_criterion = WithLossCell(net, criterion)
|
||||
train_network = TrainOneStepCell(net_with_criterion, optimizer) # optimizer
|
||||
train_network.set_train()
|
||||
res = train_network(data, label)
|
||||
print("+++++++++Loss+++++++++++++")
|
||||
print(res)
|
||||
print("+++++++++++++++++++++++++++")
|
||||
diff = res.asnumpy()[0] - 2.3025851
|
||||
assert np.all(diff < 1.e-7)
|
||||
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
def test_lenet():
|
||||
data = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32) * 0.01)
|
||||
label = Tensor(np.ones([32]).astype(np.int32))
|
||||
net = LeNet()
|
||||
train(net, data, label)
|
||||
@ -0,0 +1,174 @@
|
||||
/**
|
||||
* Copyright 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 "common/backend_common_test.h"
|
||||
#include "debug/anf_ir_dump.h"
|
||||
#include "common/py_func_graph_fetcher.h"
|
||||
#include "backend/optimizer/ascend/format_type/remove_internal_output.h"
|
||||
|
||||
#define private public
|
||||
#define protected public
|
||||
#include "backend/optimizer/ascend/format_type/insert_trans_op.h"
|
||||
#undef private
|
||||
#undef protected
|
||||
|
||||
namespace mindspore {
|
||||
namespace opt {
|
||||
using KernelBuildInfoBuilder = kernel::KernelBuildInfo::KernelBuildInfoBuilder;
|
||||
|
||||
class TestHWRemoveInternalOutput : public BackendCommon {
|
||||
public:
|
||||
TestHWRemoveInternalOutput() : getPyFun_("gtest_input.pre_activate.remove_internal_output_test", true) {}
|
||||
~TestHWRemoveInternalOutput() override = default;
|
||||
|
||||
AnfNodePtr GetMakeTuple(const KernelGraphPtr &kg) {
|
||||
auto ret = kg->get_return();
|
||||
MS_EXCEPTION_IF_NULL(ret);
|
||||
auto make_tuple = ret->input(1);
|
||||
return make_tuple;
|
||||
}
|
||||
|
||||
KernelGraphPtr GetSingleOutputGraph(const std::string &func_name, const std::string &sub_func_name) {
|
||||
FuncGraphPtr g = getPyFun_.CallAndParseRet(func_name, sub_func_name);
|
||||
std::vector<int> shp{2, 32, 224, 224};
|
||||
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
|
||||
AbstractBasePtrList args_spec_list{x_abstract, x_abstract};
|
||||
auto kg = GetKernelGraph(g, args_spec_list);
|
||||
auto make_tuple = GetMakeTuple(kg);
|
||||
auto add = make_tuple->cast<CNodePtr>()->input(1);
|
||||
MS_EXCEPTION_IF_NULL(add);
|
||||
kg->AddInternalOutput(add, add);
|
||||
KernelBuildInfoBuilder builder;
|
||||
builder.SetInputsFormat({kOpFormat_DEFAULT, kOpFormat_DEFAULT});
|
||||
builder.SetInputsDeviceType({kFloat32->type_id(), kFloat32->type_id()});
|
||||
builder.SetOutputsFormat({kOpFormat_NC1HWC0});
|
||||
builder.SetOutputsDeviceType({kFloat16->type_id()});
|
||||
add->set_kernel_info(std::make_shared<device::KernelInfo>());
|
||||
AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), add.get());
|
||||
return kg;
|
||||
}
|
||||
|
||||
KernelGraphPtr GetMutilpleOutputGraph(const std::string &func_name, const std::string &sub_func_name) {
|
||||
FuncGraphPtr g = getPyFun_.CallAndParseRet(func_name, sub_func_name);
|
||||
std::vector<int> shp{2, 32, 224, 224};
|
||||
auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
|
||||
AbstractBasePtrList args_spec_list{x_abstract};
|
||||
auto kg = GetKernelGraph(g, args_spec_list);
|
||||
auto output_make_tuple = GetMakeTuple(kg);
|
||||
auto make_tuple = output_make_tuple->cast<CNodePtr>()->input(1);
|
||||
MS_EXCEPTION_IF_NULL(make_tuple);
|
||||
auto tuple_getitem1 = make_tuple->cast<CNodePtr>()->input(1);
|
||||
MS_EXCEPTION_IF_NULL(tuple_getitem1);
|
||||
auto tuple_getitem2 = make_tuple->cast<CNodePtr>()->input(2);
|
||||
MS_EXCEPTION_IF_NULL(tuple_getitem2);
|
||||
auto max_pool = tuple_getitem1->cast<CNodePtr>()->input(1);
|
||||
MS_EXCEPTION_IF_NULL(max_pool);
|
||||
kg->AddInternalOutput(tuple_getitem1, max_pool);
|
||||
kg->AddInternalOutput(tuple_getitem2, max_pool);
|
||||
KernelBuildInfoBuilder builder;
|
||||
builder.SetInputsFormat({kOpFormat_DEFAULT});
|
||||
builder.SetInputsDeviceType({kFloat32->type_id()});
|
||||
builder.SetOutputsFormat({kOpFormat_NC1HWC0, kOpFormat_NC1HWC0});
|
||||
builder.SetOutputsDeviceType({kFloat16->type_id(), kFloat16->type_id()});
|
||||
max_pool->set_kernel_info(std::make_shared<device::KernelInfo>());
|
||||
AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), max_pool.get());
|
||||
return kg;
|
||||
}
|
||||
UT::PyFuncGraphFetcher getPyFun_;
|
||||
};
|
||||
|
||||
class MockRemoveInternalOutputTransOpKernelSelect : public KernelSelect {
|
||||
public:
|
||||
MockRemoveInternalOutputTransOpKernelSelect() = default;
|
||||
~MockRemoveInternalOutputTransOpKernelSelect() override = default;
|
||||
void SelectKernel(const CNodePtr &cnode) override {
|
||||
KernelBuildInfoBuilder builder;
|
||||
builder.SetInputsFormat({kOpFormat_NC1HWC0});
|
||||
builder.SetInputsDeviceType({kFloat16->type_id()});
|
||||
builder.SetOutputsFormat({kOpFormat_DEFAULT});
|
||||
builder.SetOutputsDeviceType({kFloat32->type_id()});
|
||||
AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), cnode.get());
|
||||
}
|
||||
};
|
||||
|
||||
TEST_F(TestHWRemoveInternalOutput, test_remove_internal_output_trans_op_for_single_output) {
|
||||
auto ms_context = MsContext::GetInstance();
|
||||
MS_EXCEPTION_IF_NULL(ms_context);
|
||||
ms_context->set_execution_mode(kGraphMode);
|
||||
auto kg = GetSingleOutputGraph("test_remove_internal_output_trans_op_for_single_output", "before");
|
||||
// insert trans op for output
|
||||
auto graph_optimizer = std::make_shared<opt::GraphOptimizer>();
|
||||
auto pass_manager = std::make_shared<opt::PassManager>();
|
||||
auto insert_trans_op_pass = std::make_shared<opt::InsertTransOp>();
|
||||
insert_trans_op_pass->kernel_select_ = std::make_shared<MockRemoveInternalOutputTransOpKernelSelect>();
|
||||
pass_manager->AddPass(insert_trans_op_pass);
|
||||
graph_optimizer->AddPassManager(pass_manager);
|
||||
auto new_g = graph_optimizer->Optimize(kg);
|
||||
FuncGraphPtr g_after =
|
||||
getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_single_output", "after_insert_trans_op");
|
||||
EXPECT_TRUE(CheckEqualGraph(g_after, new_g));
|
||||
|
||||
auto make_tuple = GetMakeTuple(kg);
|
||||
auto trans_data = make_tuple->cast<CNodePtr>()->input(1);
|
||||
EXPECT_TRUE(kg->IsInternalOutput(trans_data));
|
||||
|
||||
// remove trans op for internal output
|
||||
auto graph_optimizer1 = std::make_shared<opt::GraphOptimizer>();
|
||||
auto pass_manager1 = std::make_shared<opt::PassManager>();
|
||||
auto remove_internal_output_trans_op_pass = std::make_shared<opt::RemoveInternalOutputTransOp>();
|
||||
pass_manager1->AddPass(remove_internal_output_trans_op_pass);
|
||||
graph_optimizer1->AddPassManager(pass_manager1);
|
||||
auto new_g1 = graph_optimizer1->Optimize(new_g);
|
||||
FuncGraphPtr g_after1 = getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_single_output",
|
||||
"after_remove_internal_output_trans_op");
|
||||
EXPECT_TRUE(CheckEqualGraph(g_after1, new_g1));
|
||||
}
|
||||
|
||||
TEST_F(TestHWRemoveInternalOutput, test_remove_internal_output_trans_op_for_multiple_output) {
|
||||
auto kg = GetMutilpleOutputGraph("test_remove_internal_output_trans_op_for_multiple_output", "before");
|
||||
// insert trans op for output
|
||||
auto graph_optimizer = std::make_shared<opt::GraphOptimizer>();
|
||||
auto pass_manager = std::make_shared<opt::PassManager>();
|
||||
auto insert_trans_op_pass = std::make_shared<opt::InsertTransOp>();
|
||||
insert_trans_op_pass->kernel_select_ = std::make_shared<MockRemoveInternalOutputTransOpKernelSelect>();
|
||||
pass_manager->AddPass(insert_trans_op_pass);
|
||||
graph_optimizer->AddPassManager(pass_manager);
|
||||
auto new_g = graph_optimizer->Optimize(kg);
|
||||
FuncGraphPtr g_after =
|
||||
getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_multiple_output", "after_insert_trans_op");
|
||||
EXPECT_TRUE(CheckEqualGraph(g_after, new_g));
|
||||
|
||||
auto output_make_tuple = GetMakeTuple(kg);
|
||||
auto make_tuple = output_make_tuple->cast<CNodePtr>()->input(1);
|
||||
auto tuple_getitem = make_tuple->cast<CNodePtr>()->input(1);
|
||||
auto make_tuple1 = tuple_getitem->cast<CNodePtr>()->input(1);
|
||||
auto trans_data1 = make_tuple1->cast<CNodePtr>()->input(1);
|
||||
auto trans_data2 = make_tuple1->cast<CNodePtr>()->input(2);
|
||||
EXPECT_TRUE(kg->IsInternalOutput(trans_data1));
|
||||
EXPECT_TRUE(kg->IsInternalOutput(trans_data2));
|
||||
|
||||
// remove trans op for internal output
|
||||
auto graph_optimizer1 = std::make_shared<opt::GraphOptimizer>();
|
||||
auto pass_manager1 = std::make_shared<opt::PassManager>();
|
||||
auto remove_internal_output_trans_op_pass = std::make_shared<opt::RemoveInternalOutputTransOp>();
|
||||
pass_manager1->AddPass(remove_internal_output_trans_op_pass);
|
||||
graph_optimizer1->AddPassManager(pass_manager1);
|
||||
auto new_g1 = graph_optimizer1->Optimize(new_g);
|
||||
FuncGraphPtr g_after1 = getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_multiple_output",
|
||||
"after_remove_internal_output_trans_op");
|
||||
EXPECT_TRUE(CheckEqualGraph(g_after1, new_g1));
|
||||
}
|
||||
} // namespace opt
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,83 @@
|
||||
# Copyright 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.
|
||||
# ============================================================================
|
||||
from mindspore.ops import Primitive
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
tuple_getitem = Primitive('tuple_getitem')
|
||||
add = P.TensorAdd()
|
||||
max_pool = P.MaxPoolWithArgmax(padding="same", ksize=3, strides=2)
|
||||
make_tuple = Primitive('make_tuple')
|
||||
trans_data = Primitive("TransData")
|
||||
|
||||
|
||||
class FnDict:
|
||||
def __init__(self):
|
||||
self.fnDict = {}
|
||||
|
||||
def __call__(self, fn):
|
||||
self.fnDict[fn.__name__] = fn
|
||||
|
||||
def __getitem__(self, name):
|
||||
return self.fnDict[name]
|
||||
|
||||
|
||||
def test_remove_internal_output_trans_op_for_single_output(tag):
|
||||
fns = FnDict()
|
||||
|
||||
@fns
|
||||
def before(x, y):
|
||||
res = add(x, y)
|
||||
return res
|
||||
|
||||
@fns
|
||||
def after_insert_trans_op(x, y):
|
||||
output = add(x, y)
|
||||
res = trans_data(output)
|
||||
return make_tuple(res)
|
||||
|
||||
@fns
|
||||
def after_remove_internal_output_trans_op(x, y):
|
||||
res = add(x, y)
|
||||
return make_tuple(res)
|
||||
|
||||
return fns[tag]
|
||||
|
||||
|
||||
def test_remove_internal_output_trans_op_for_multiple_output(tag):
|
||||
fns = FnDict()
|
||||
|
||||
@fns
|
||||
def before(x):
|
||||
max_pool_res = max_pool(x)
|
||||
res = make_tuple(tuple_getitem(max_pool_res, 0), tuple_getitem(max_pool_res, 1))
|
||||
return res
|
||||
|
||||
@fns
|
||||
def after_insert_trans_op(x):
|
||||
output = max_pool(x)
|
||||
trans_data0 = trans_data(tuple_getitem(output, 0))
|
||||
trans_data1 = trans_data(tuple_getitem(output, 1))
|
||||
new_make_tuple = make_tuple(trans_data0, trans_data1)
|
||||
res = make_tuple(tuple_getitem(new_make_tuple, 0), tuple_getitem(new_make_tuple, 1))
|
||||
return make_tuple(res)
|
||||
|
||||
@fns
|
||||
def after_remove_internal_output_trans_op(x):
|
||||
output = max_pool(x)
|
||||
new_make_tuple = make_tuple(tuple_getitem(output, 0), tuple_getitem(output, 1))
|
||||
res = make_tuple(tuple_getitem(new_make_tuple, 0), tuple_getitem(new_make_tuple, 1))
|
||||
return make_tuple(res)
|
||||
|
||||
return fns[tag]
|
||||
Loading…
Reference in new issue