From 1dcf6a72126c4da53f46b180a63a46fa937ffd5a Mon Sep 17 00:00:00 2001 From: Huihuang Zheng Date: Fri, 6 Dec 2019 18:08:43 +0800 Subject: [PATCH] Add Much Complex Test and Fix Bugs for Control Flow cond API (#21532) Add tests to use dy/dx to make sure the gradient values calculated by the control flow backward is correct. Also fixed bugs detected by those tests. Fix bugs: 1. Unlike sum_op, optimizer ops don't allow uninitialized input tensor. But in conditional_block_grad_op, since the conditional_block may not run, the output gradient tensor may be uninitialized, which will cause the optimizer op error. To fix it, we should let optimizer ops support uninitialized input like sum_op or assign the uninitialized gradient to 0 when the conditional_block_grad_op doesn't run. I found there are about 10+ optimizer ops. **To be simpler, I just assign output gradient of the conditional_block_grad_op to 0 in this PR**. But it can be further explored whether we can make optimizer ops like sum_op to support uninitialized input tensor because theoretically we can speed up without the assigning in conditional_block_grad_op. 2. Infer parameter shapes during append_backward. I didn't know that all our parameters are in global block. When op_desc is inferring shapes at the sub-block, it may not know the shape of gradients of parameters whose shape information is at global block. I fixed it by inferring shapes of gradients from forward var. This PR also did some code clean up: 1. Print the var name when sgd_op catches shape error so that it is easier to debug 2. Fix a typo: dicta -> dict --- .../operators/controlflow/CMakeLists.txt | 2 + .../controlflow/conditional_block_op.cc | 83 ++++++++++- .../controlflow/conditional_block_op_test.cc | 80 +++++++++++ paddle/fluid/operators/optimizers/sgd_op.cc | 5 +- python/paddle/fluid/backward.py | 10 +- python/paddle/fluid/layers/control_flow.py | 7 +- .../fluid/tests/unittests/simple_nets.py | 13 +- .../paddle/fluid/tests/unittests/test_cond.py | 132 ++++++++++++++++++ .../unittests/test_fuse_all_reduce_pass.py | 10 +- .../unittests/test_select_input_output_op.py | 65 ++++++--- 10 files changed, 367 insertions(+), 40 deletions(-) create mode 100644 paddle/fluid/operators/controlflow/conditional_block_op_test.cc diff --git a/paddle/fluid/operators/controlflow/CMakeLists.txt b/paddle/fluid/operators/controlflow/CMakeLists.txt index 758f0a65d1..7ae54ef480 100644 --- a/paddle/fluid/operators/controlflow/CMakeLists.txt +++ b/paddle/fluid/operators/controlflow/CMakeLists.txt @@ -5,6 +5,8 @@ cc_library(conditional_block_op_helper SRCS conditional_block_op_helper.cc DEPS cc_library(recurrent_op_helper SRCS recurrent_op_helper.cc DEPS operator op_variant recurrent_op) cc_library(while_op_helper SRCS while_op_helper.cc DEPS operator op_variant) +cc_test(conditional_block_op_test SRCS conditional_block_op_test.cc DEPS conditional_block_op executor) + target_link_libraries(conditional_block_infer_op conditional_block_op) file(APPEND ${pybind_file} "USE_OP(less_than);\nUSE_OP(logical_and);\nUSE_NO_KERNEL_OP(read_from_array);\n") diff --git a/paddle/fluid/operators/controlflow/conditional_block_op.cc b/paddle/fluid/operators/controlflow/conditional_block_op.cc index c74c4ebbd8..f22f584213 100644 --- a/paddle/fluid/operators/controlflow/conditional_block_op.cc +++ b/paddle/fluid/operators/controlflow/conditional_block_op.cc @@ -13,7 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/fluid/operators/controlflow/conditional_block_op.h" + #include "paddle/fluid/operators/assign_op.h" +#include "paddle/fluid/operators/math/math_function.h" namespace paddle { namespace operators { @@ -94,11 +96,10 @@ class ConditionalBlockGradOp : public ConditionalOp { [](const framework::LoDTensor *t) { return t->numel() != 0; }); } + const auto &inputs = Inputs(ConditionalOp::kInputs); + const auto &outside_grads = + Outputs(framework::GradVarName(ConditionalOp::kInputs)); if (need_run) { - const auto &inputs = Inputs(ConditionalOp::kInputs); - const auto &outside_grads = - Outputs(framework::GradVarName(ConditionalOp::kInputs)); - std::vector inside_grads; inside_grads.reserve(inputs.size()); for (auto &in : inputs) { @@ -126,7 +127,10 @@ class ConditionalBlockGradOp : public ConditionalOp { AssignLocalGradientToParentScope(dev_place, cur_scope, scope, inside_grads, outside_grads); + return; } + + AssignZeroToParentScope(dev_place, scope, inputs, outside_grads); } private: @@ -156,6 +160,77 @@ class ConditionalBlockGradOp : public ConditionalOp { AssignFunctor(outside_var, *dev_ctx)); } } + + void AssignZeroToParentScope( + const platform::Place &place, const framework::Scope &scope, + const std::vector &inputs, + const std::vector &outside_grads) const { + for (size_t i = 0; i < outside_grads.size(); ++i) { + const std::string &outside_grad_name = outside_grads[i]; + const std::string &input_name = inputs[i]; + VLOG(4) << "input_name = " << input_name + << ", outside_grad_name = " << outside_grad_name; + framework::Variable *input_var = scope.FindVar(input_name); + if (input_var == nullptr) { + continue; + } + framework::Variable *outside_var = scope.FindVar(outside_grad_name); + if (outside_var == nullptr) { + continue; + } + + if (input_var->IsType()) { + PADDLE_ENFORCE_EQ(outside_var->IsType(), true, + platform::errors::InvalidArgument( + "Type of outside_var %s is NOT LoDTensor, which " + "doesn't match input_var %s", + outside_grad_name, input_name)); + AssignZeroToOutsideTensor( + place, scope, input_var->Get(), + outside_var->GetMutable()); + } else if (input_var->IsType()) { + PADDLE_ENFORCE_EQ(outside_var->IsType(), + true, + platform::errors::InvalidArgument( + "Type of outside_var %s is NOT LoDTensorArray, " + "which doesn't match input_var %s", + outside_grad_name, input_name)); + const auto &input_tensors = input_var->Get(); + auto *outside_tensors = + outside_var->GetMutable(); + PADDLE_ENFORCE_EQ(input_tensors.size(), outside_tensors->size(), + platform::errors::InvalidArgument( + "LoDTensorArray outside_var %s doen't have same " + "size as input_var %s", + outside_grad_name, input_name)); + for (size_t j = 0; j < input_tensors.size(); ++j) { + AssignZeroToOutsideTensor(place, scope, input_tensors[j], + &((*outside_tensors)[j])); + } + } else { + // TODO(huihuangzheng): add support for SelectedRows + PADDLE_THROW(platform::errors::InvalidArgument( + "Conditional block grad op doesn't support non-LoDTensor output " + "now")); + } + } + } + + void AssignZeroToOutsideTensor(const platform::Place &place, + const framework::Scope &cur_scope, + const framework::LoDTensor &input_tensor, + framework::LoDTensor *outside_tensor) const { + if (!input_tensor.IsInitialized() || input_tensor.numel() == 0) { + return; + } + VLOG(4) << "Assigning zero to " << outside_tensor; + outside_tensor->Resize(input_tensor.dims()); + outside_tensor->mutable_data(place, input_tensor.type()); + const platform::DeviceContext *dev_ctx = + platform::DeviceContextPool::Instance().Get(place); + math::set_constant(*dev_ctx, outside_tensor, 0.0f); + outside_tensor->set_lod(input_tensor.lod()); + } }; class ConditionalBlockGradInferShape : public framework::InferShapeBase { diff --git a/paddle/fluid/operators/controlflow/conditional_block_op_test.cc b/paddle/fluid/operators/controlflow/conditional_block_op_test.cc new file mode 100644 index 0000000000..a5ca4a289d --- /dev/null +++ b/paddle/fluid/operators/controlflow/conditional_block_op_test.cc @@ -0,0 +1,80 @@ +/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. + +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 "paddle/fluid/operators/controlflow/conditional_block_op.h" +#include +#include +#include +#include "gtest/gtest.h" +#include "paddle/fluid/framework/executor.h" +#include "paddle/fluid/framework/op_registry.h" +#include "paddle/fluid/framework/scope.h" +#include "paddle/fluid/framework/var_type.h" + +USE_NO_KERNEL_OP(conditional_block); +USE_NO_KERNEL_OP(conditional_block_grad); + +using LoDTensor = paddle::framework::LoDTensor; +using LoDTensorArray = paddle::framework::LoDTensorArray; +using Scope = paddle::framework::Scope; +using Variable = paddle::framework::Variable; +using Place = paddle::platform::Place; + +TEST(ConditionalBlockGrad, NoNeedRunLoDTensorArray) { + Place place = paddle::platform::CPUPlace(); + Scope scope; + + Variable* cond_var = scope.Var("condition"); + LoDTensor* cond_tensor = cond_var->GetMutable(); + paddle::framework::DDim cond_dims = paddle::framework::make_ddim({1}); + bool* cond_data = cond_tensor->mutable_data(cond_dims, place); + cond_data[0] = false; + + Variable* input_var = scope.Var("input_lod_tensor_array"); + LoDTensorArray* input_tensors = input_var->GetMutable(); + for (int i = 0; i < 5; ++i) { + paddle::framework::DDim in_dims = + paddle::framework::make_ddim({i + 1, i + 2}); + LoDTensor lod_tensor; + float* in_data = lod_tensor.mutable_data(in_dims, place); + for (int j = 0; j < (i + 1) * (i + 2); ++j) { + in_data[j] = static_cast(j); + } + input_tensors->push_back(lod_tensor); + } + + Variable* input_grad_var = scope.Var("input_lod_tensor_array@GRAD"); + LoDTensorArray* grad_tensors = input_grad_var->GetMutable(); + grad_tensors->resize(5); + + paddle::framework::AttributeMap attrs; + attrs.insert({"is_scalar_condition", true}); + + auto conditional_grad_op = paddle::framework::OpRegistry::CreateOp( + "conditional_block_grad", + {{"Input", {"input_lod_tensor_array"}}, {"Cond", {"condition"}}}, + {{"Input@GRAD", {"input_lod_tensor_array@GRAD"}}}, attrs); + + conditional_grad_op->Run(scope, place); + + const LoDTensorArray& out_tensors = input_grad_var->Get(); + for (int i = 0; i < 5; ++i) { + paddle::framework::DDim out_dims = out_tensors[i].dims(); + EXPECT_EQ(paddle::framework::make_ddim({i + 1, i + 2}), out_dims); + const float* out_data = out_tensors[i].data(); + for (int j = 0; j < (i + 1) * (i + 2); ++j) { + EXPECT_EQ(0, out_data[j]); + } + } +} diff --git a/paddle/fluid/operators/optimizers/sgd_op.cc b/paddle/fluid/operators/optimizers/sgd_op.cc index 89e8d3e141..14db2c08d1 100644 --- a/paddle/fluid/operators/optimizers/sgd_op.cc +++ b/paddle/fluid/operators/optimizers/sgd_op.cc @@ -46,8 +46,9 @@ class SGDOp : public framework::OperatorWithKernel { param_dim, ctx->GetInputDim("Grad"), platform::errors::InvalidArgument( "SGD Operator's input Param and Grad dimensions do not match. " - "The Param shape is [%s], but the Grad shape is [%s].", - param_dim, ctx->GetInputDim("Grad"))); + "The Param %s shape is [%s], but the Grad %s shape is [%s].", + ctx->Inputs("Param")[0], param_dim, ctx->Inputs("Grad")[0], + ctx->GetInputDim("Grad"))); } ctx->SetOutputDim("ParamOut", param_dim); } diff --git a/python/paddle/fluid/backward.py b/python/paddle/fluid/backward.py index 54e7066ba5..0d2c7e31fb 100644 --- a/python/paddle/fluid/backward.py +++ b/python/paddle/fluid/backward.py @@ -263,15 +263,16 @@ def _create_loss_op_desc_(loss): return op_desc -def _infer_var_data_type_(grad_var_name, block): +def _infer_var_data_type_shape_(grad_var_name, block): """ - Infer the data type of given grad variable + Infer the data type and shape of given grad variable """ grad_var = block.desc.find_var(cpt.to_bytes(grad_var_name)) fwd_name = _strip_grad_suffix_(grad_var_name) if block.desc.has_var_recursive(cpt.to_bytes(fwd_name)): fwd_var = block.desc.find_var_recursive(cpt.to_bytes(fwd_name)) grad_var.set_dtype(fwd_var.dtype()) + grad_var.set_shape(fwd_var.shape()) else: grad_var.set_dtype(core.VarDesc.VarType.FP32) @@ -921,9 +922,10 @@ def _append_backward_vars_(block, start_op_idx, grad_to_var, grad_info_map): # infer_shape and infer_type op_desc.infer_var_type(block.desc) op_desc.infer_shape(block.desc) + for arg in op_desc.output_arg_names(): if arg in new_vars: - _infer_var_data_type_(arg, block) + _infer_var_data_type_shape_(arg, block) def _rename_grad_(block, start_op_idx, grad_to_var, target_grad_map): @@ -1062,7 +1064,6 @@ def append_backward(loss, no_grad_dict = _get_stop_gradients_(program) no_grad_dict[0].update(list(map(_append_grad_suffix_, no_grad_set))) - grad_info_map = dict() root_block = program.block(0) fwd_op_num = root_block.desc.op_size() @@ -1114,6 +1115,7 @@ def append_backward(loss, # different names. _rename_grad_(root_block, fwd_op_num, grad_to_var, {}) + grad_info_map = dict() _append_backward_vars_(root_block, fwd_op_num, grad_to_var, grad_info_map) program.current_block_idx = current_block_idx diff --git a/python/paddle/fluid/layers/control_flow.py b/python/paddle/fluid/layers/control_flow.py index 4d42877ad4..edae54294b 100755 --- a/python/paddle/fluid/layers/control_flow.py +++ b/python/paddle/fluid/layers/control_flow.py @@ -80,9 +80,11 @@ def select_input(inputs, mask): helper = LayerHelper('select_input', **locals()) if isinstance(inputs, list) or isinstance(inputs, tuple): input_dtype = inputs[0].dtype + input_shape = inputs[0].shape else: input_dtype = inputs.dtype - out = helper.create_variable(dtype=input_dtype) + input_shape = inputs.shape + out = helper.create_variable(dtype=input_dtype, shape=input_shape) helper.append_op( type='select_input', inputs={'X': inputs, @@ -1742,7 +1744,8 @@ def copy_var_to_parent_block(var, layer_helper): assert parent_idx >= 0, "Got wrong parent block index when assigning var to parent scope in control_flow" parent_block = prog.block(parent_idx) - parent_block_var = parent_block.create_var(dtype=var.dtype, type=var.type) + parent_block_var = parent_block.create_var( + dtype=var.dtype, shape=var.shape, type=var.type) assign(var, parent_block_var) return parent_block_var diff --git a/python/paddle/fluid/tests/unittests/simple_nets.py b/python/paddle/fluid/tests/unittests/simple_nets.py index f65bc7c3f4..5f48cc8e6a 100644 --- a/python/paddle/fluid/tests/unittests/simple_nets.py +++ b/python/paddle/fluid/tests/unittests/simple_nets.py @@ -37,10 +37,7 @@ def simple_fc_net(use_feed=None): return simple_fc_net_with_inputs(img, label, class_num=10) -def fc_with_batchnorm(use_feed=None): - img = fluid.layers.data(name='image', shape=[784], dtype='float32') - label = fluid.layers.data(name='label', shape=[1], dtype='int64') - +def batchnorm_fc_with_inputs(img, label, class_num=10): hidden = img for _ in range(2): hidden = fluid.layers.fc( @@ -52,12 +49,18 @@ def fc_with_batchnorm(use_feed=None): hidden = fluid.layers.batch_norm(input=hidden) - prediction = fluid.layers.fc(hidden, size=10, act='softmax') + prediction = fluid.layers.fc(hidden, size=class_num, act='softmax') loss = fluid.layers.cross_entropy(input=prediction, label=label) loss = fluid.layers.mean(loss) return loss +def fc_with_batchnorm(use_feed=None): + img = fluid.layers.data(name='image', shape=[784], dtype='float32') + label = fluid.layers.data(name='label', shape=[1], dtype='int64') + return batchnorm_fc_with_inputs(img, label, class_num=10) + + def bow_net(use_feed, dict_dim, is_sparse=False, diff --git a/python/paddle/fluid/tests/unittests/test_cond.py b/python/paddle/fluid/tests/unittests/test_cond.py index acbe97dd12..f7df7c83d2 100644 --- a/python/paddle/fluid/tests/unittests/test_cond.py +++ b/python/paddle/fluid/tests/unittests/test_cond.py @@ -24,6 +24,9 @@ import paddle.fluid.framework as framework from paddle.fluid.backward import append_backward from paddle.fluid.executor import Executor from paddle.fluid.framework import Program, program_guard +from simple_nets import simple_fc_net_with_inputs, batchnorm_fc_with_inputs + +np.random.seed(123) class TestCondInputOutput(unittest.TestCase): @@ -275,5 +278,134 @@ class TestCondNestedControlFlow(unittest.TestCase): self.assertEqual(ret[1][0], expected_a_grad) +class TestCondBackward(unittest.TestCase): + def backward_value_helper(self, cond_func): + """ + Helper function that compares calculated backward value is close to dy/dx + """ + main_program = Program() + main_program.random_seed = 123 + startup_program = Program() + startup_program.random_seed = 123 + with program_guard(main_program, startup_program): + img = fluid.data(name='image', shape=[-1, 9], dtype='float32') + img.stop_gradient = False + label = fluid.data(name='label', shape=[-1, 1], dtype='int64') + i = fluid.data(name="i", shape=[1], dtype='int32') + loss = cond_func(i, img, label) + append_backward(loss) + place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda( + ) else fluid.CPUPlace() + exe = fluid.Executor(place) + exe.run(startup_program) + + delta = 0.005 + for feed_i in range(0, 10): + feed_img = np.random.random(size=[1, 9]).astype(np.float32) + feed_label = np.random.randint( + low=0, high=10, size=[1, 1], dtype=np.int64) + img_grad, loss_value = exe.run( + main_program, + feed={ + 'i': np.full((1), feed_i, np.int32), + 'image': feed_img, + 'label': feed_label + }, + fetch_list=[img.grad_name, loss.name]) + + numerical_grad = np.zeros(shape=[1, 9], dtype=np.float32) + feed_img_delta = np.copy(feed_img) + for j in range(9): + feed_img_delta[0][j] = feed_img[0][j] + delta + loss_delta = exe.run(main_program, + feed={ + 'i': np.full((1), feed_i, np.int32), + 'image': feed_img_delta, + 'label': feed_label + }, + fetch_list=[loss.name]) + numerical_grad[0][j] = (loss_delta[0] - loss_value[0]) / delta + feed_img_delta[0][j] = feed_img[0][j] + self.assertTrue( + np.isclose( + img_grad, numerical_grad, atol=0.05, rtol=0.05).all()) + + def add_optimizer_helper(self, cond_func): + """ + Test that program is runnable when add optimizer + """ + main_program = Program() + startup_program = Program() + with program_guard(main_program, startup_program): + img = fluid.data(name='image', shape=[-1, 784], dtype='float32') + label = fluid.data(name='label', shape=[-1, 1], dtype='int64') + i = fluid.data(name="i", shape=[1], dtype='int32') + loss = cond_func(i, img, label) + optimizer = fluid.optimizer.SGD(learning_rate=0.1) + optimizer.minimize(loss) + + place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda( + ) else fluid.CPUPlace() + exe = fluid.Executor(place) + exe.run(startup_program) + + for feed_i in range(0, 10): + feed_img = np.random.random(size=[16, 784]).astype(np.float32) + feed_label = np.random.randint( + low=0, high=10, size=[16, 1], dtype=np.int64) + exe.run(main_program, + feed={ + 'i': np.full((1), feed_i, np.int32), + 'image': feed_img, + 'label': feed_label + }, + fetch_list=[loss]) + + def test_cond_backward(self): + def cond_func(i, img, label): + predicate = ((i % 2) == 0) + return layers.cond(predicate, + lambda: simple_fc_net_with_inputs(img, label, class_num=10), + lambda: batchnorm_fc_with_inputs(img, label, class_num=10)) + + self.backward_value_helper(cond_func) + self.add_optimizer_helper(cond_func) + + def test_half_nested_cond_backward(self): + def branch(i, img, label): + return layers.cond((i % 2) == 0, lambda: simple_fc_net_with_inputs(img, label, class_num=10), + lambda: batchnorm_fc_with_inputs(img, label, class_num=10)) + + def cond_func_simple_net_at_true(i, img, label): + return layers.cond(i < 5, lambda: branch(i, img, label), + lambda: layers.mean(img)) + + def cond_func_simple_net_at_false(i, img, label): + return layers.cond(i < 5, lambda: layers.mean(img), + lambda: branch(i, img, label)) + + self.backward_value_helper(cond_func_simple_net_at_true) + self.add_optimizer_helper(cond_func_simple_net_at_true) + self.backward_value_helper(cond_func_simple_net_at_false) + self.add_optimizer_helper(cond_func_simple_net_at_false) + + def test_nested_cond_backward(self): + def branch(i, img, label, mod_two): + + if mod_two: + predicate = ((i % 2) == 0) + else: + predicate = ((i % 2) != 0) + return layers.cond(predicate, lambda: simple_fc_net_with_inputs(img, label, class_num=10), + lambda: batchnorm_fc_with_inputs(img, label, class_num=10)) + + def cond_func(i, img, label): + return layers.cond(i < 5, lambda: branch(i, img, label, True), + lambda: branch(i, img, label, False)) + + self.backward_value_helper(cond_func) + self.add_optimizer_helper(cond_func) + + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/fluid/tests/unittests/test_fuse_all_reduce_pass.py b/python/paddle/fluid/tests/unittests/test_fuse_all_reduce_pass.py index 5ce82b267a..06f8da84a2 100644 --- a/python/paddle/fluid/tests/unittests/test_fuse_all_reduce_pass.py +++ b/python/paddle/fluid/tests/unittests/test_fuse_all_reduce_pass.py @@ -30,7 +30,7 @@ class TestFuseAllReduceOpsBase(TestParallelExecutorBase): def compare_fuse_all_reduce_ops(self, model, use_cuda, - init_feed_dicta=None, + init_feed_dict=None, get_data_from_feeder=None, optimizer=None, fuse_all_optimizer_ops=False): @@ -38,8 +38,8 @@ class TestFuseAllReduceOpsBase(TestParallelExecutorBase): return feed_dict_data = None - if init_feed_dicta is not None: - img, label = init_feed_dicta() + if init_feed_dict is not None: + img, label = init_feed_dict() feed_dict_data = {"image": img, "label": label} not_fuse_op_first_loss, not_fuse_op_last_loss = self.check_network_convergence( @@ -76,7 +76,7 @@ class TestFuseAllReduceOps(TestFuseAllReduceOpsBase): self.compare_fuse_all_reduce_ops( model, use_cuda, - init_feed_dicta=init_data, + init_feed_dict=init_data, optimizer=self.optimizer, fuse_all_optimizer_ops=True) @@ -94,7 +94,7 @@ class TestFuseAllReduceOpsAndOptiOps(TestFuseAllReduceOps): self.compare_fuse_all_reduce_ops( model, use_cuda, - init_feed_dicta=init_data, + init_feed_dict=init_data, optimizer=self.optimizer, fuse_all_optimizer_ops=True) diff --git a/python/paddle/fluid/tests/unittests/test_select_input_output_op.py b/python/paddle/fluid/tests/unittests/test_select_input_output_op.py index 092262da47..bd66bf81b0 100644 --- a/python/paddle/fluid/tests/unittests/test_select_input_output_op.py +++ b/python/paddle/fluid/tests/unittests/test_select_input_output_op.py @@ -26,22 +26,52 @@ from paddle.fluid.layers.control_flow import select_input, select_output class TestSplitMergeSelectedVarOps(unittest.TestCase): - def test_forward_backward(self): - branch_num = 9 + def test_forward_backward_list_output(self): + for branch_num in range(2, 10): + program = Program() + with program_guard(program): + x = layers.data(name='x', shape=[2], dtype='float32') + x.stop_gradient = False # For test gradient + mask = layers.data(name='mask', shape=[1], dtype='int32') + + outputs = [] + for i in range(branch_num): + out = program.current_block().create_var( + dtype='float32', type=core.VarDesc.VarType.LOD_TENSOR) + outputs.append(out) + + select_output(x, outputs, mask) + y = select_input(outputs, mask) + mean = layers.mean(y) + append_backward(mean) + + place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda( + ) else fluid.CPUPlace() + exe = Executor(place) + + feed_x = np.asarray([1.3, -1.4]).astype(np.float32) + for i in range(branch_num): + feed_mask = np.asarray([i]).astype(np.int32) + ret = exe.run(program, + feed={'x': feed_x, + 'mask': feed_mask}, + fetch_list=[y.name, x.grad_name]) + x_grad = np.asarray([0.5, 0.5]).astype(np.float32) + self.assertTrue(np.allclose(np.asarray(ret[0]), feed_x)) + self.assertTrue(np.allclose(np.asarray(ret[1]), x_grad)) + + def test_forward_backward_single_tensor_output(self): program = Program() with program_guard(program): x = layers.data(name='x', shape=[2], dtype='float32') x.stop_gradient = False # For test gradient mask = layers.data(name='mask', shape=[1], dtype='int32') - outputs = [] - for i in range(branch_num): - out = program.current_block().create_var( - dtype='float32', type=core.VarDesc.VarType.LOD_TENSOR) - outputs.append(out) + out = program.current_block().create_var( + dtype='float32', type=core.VarDesc.VarType.LOD_TENSOR) - select_output(x, outputs, mask) - y = select_input(outputs, mask) + select_output(x, out, mask) + y = select_input(out, mask) mean = layers.mean(y) append_backward(mean) @@ -50,15 +80,14 @@ class TestSplitMergeSelectedVarOps(unittest.TestCase): exe = Executor(place) feed_x = np.asarray([1.3, -1.4]).astype(np.float32) - for i in range(branch_num): - feed_mask = np.asarray([i]).astype(np.int32) - ret = exe.run(program, - feed={'x': feed_x, - 'mask': feed_mask}, - fetch_list=[y.name, x.grad_name]) - x_grad = np.asarray([0.5, 0.5]).astype(np.float32) - self.assertTrue(np.allclose(np.asarray(ret[0]), feed_x)) - self.assertTrue(np.allclose(np.asarray(ret[1]), x_grad)) + feed_mask = np.asarray([0]).astype(np.int32) + ret = exe.run(program, + feed={'x': feed_x, + 'mask': feed_mask}, + fetch_list=[y.name, x.grad_name]) + x_grad = np.asarray([0.5, 0.5]).astype(np.float32) + self.assertTrue(np.allclose(np.asarray(ret[0]), feed_x)) + self.assertTrue(np.allclose(np.asarray(ret[1]), x_grad)) if __name__ == '__main__': -- GitLab