Add basic Python Cond Layer (#21050)

python/paddle/fluid/layers/control_flow.py

@@ -16,12 +16,13 @@ from __future__ import print_function
 from ..wrapped_decorator import signature_safe_contextmanager
 
 from .layer_function_generator import autodoc, templatedoc
-from .tensor import assign, fill_constant
+from .tensor import assign, cast, fill_constant
 from .. import core
 from ..framework import Program, Variable, Operator
 from ..layer_helper import LayerHelper, unique_name
 from ..initializer import force_init_on_cpu
 from .nn import logical_and, logical_not, logical_or
+from .utils import assert_same_structure, flatten, map_structure
 import numpy
 import warnings
 import six
@@ -30,7 +31,7 @@ from functools import reduce
 __all__ = [
     'While', 'Switch', 'increment', 'array_write', 'create_array', 'less_than',
     'less_equal', 'greater_than', 'greater_equal', 'equal', 'not_equal',
-    'array_read', 'array_length', 'IfElse', 'DynamicRNN', 'StaticRNN',
+    'array_read', 'array_length', 'cond', 'IfElse', 'DynamicRNN', 'StaticRNN',
     'reorder_lod_tensor_by_rank', 'Print', 'is_empty'
 ]
 
@@ -1729,6 +1730,73 @@ class ConditionalBlock(object):
             })
 
 
+def copy_var_to_parent_block(var, layer_helper):
+    if var is None:
+        return None
+    prog = layer_helper.main_program
+    parent_idx = prog.current_block().parent_idx
+    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)
+    assign(var, parent_block_var)
+    return parent_block_var
+
+
+def cond(pred, true_fn=None, false_fn=None, name=None):
+    """
+    TODO:(huihuangzheng) developing
+    """
+    helper = LayerHelper('cond', **locals())
+    true_output = None
+    false_output = None
+    copy_to_global_func = lambda var: copy_var_to_parent_block(var, helper)
+    if true_fn is not None:
+        if not callable(true_fn):
+            raise TypeError("The true_fn in cond must be callable")
+        true_cond_block = ConditionalBlock([pred], is_scalar_condition=True)
+        with true_cond_block.block():
+            origin_true_output = true_fn()
+            if origin_true_output is not None:
+                true_output = map_structure(copy_to_global_func,
+                                            origin_true_output)
+    if false_fn is not None:
+        if not callable(false_fn):
+            raise TypeError("The false_fn in cond must be callable")
+        false_cond_block = ConditionalBlock(
+            [logical_not(pred)], is_scalar_condition=True)
+        with false_cond_block.block():
+            origin_false_output = false_fn()
+            if origin_false_output is not None:
+                false_output = map_structure(copy_to_global_func,
+                                             origin_false_output)
+
+    if true_output is None and false_output is None:
+        return None
+
+    if true_output is None:
+        raise ValueError(
+            "Incompatible return values of true_fn and false_fn in cond: "
+            "true_fn returns None while false_fn returns non-None")
+    if false_output is None:
+        raise ValueError(
+            "Incompatible return values of true_fn and false_fn in cond: "
+            "true_fn returns non-None while false_fn returns None")
+
+    # Merge ture and false output if they are not None
+    try:
+        assert_same_structure(true_output, false_output, check_types=False)
+    except ValueError as e:
+        raise ValueError(
+            "Incompatible return values of true_fn and false_fn in cond: {}".
+            format(e))
+
+    mask = cast(pred, dtype='int32')
+    merge_func = lambda false_var, true_var : select_input([false_var, true_var], mask)
+    merged_output = map_structure(merge_func, false_output, true_output)
+    return merged_output
+
+
 class Switch(object):
     """
 

python/paddle/fluid/tests/unittests/test_cond.py

+#   Copyright (c) 2018 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.
+
+from __future__ import print_function
+
+import numpy as np
+import unittest
+
+import paddle.fluid as fluid
+import paddle.fluid.core as core
+import paddle.fluid.layers as layers
+import paddle.fluid.framework as framework
+from paddle.fluid.executor import Executor
+from paddle.fluid.framework import Program, program_guard
+
+
+class TestCond(unittest.TestCase):
+    def test_return_single_var(self):
+        """
+        pseudocode:
+
+        if 0.23 < 0.1:
+            return 2
+        else:
+            return -1
+        """
+
+        def true_func():
+            return layers.fill_constant(shape=[2, 3], dtype='int32', value=2)
+
+        def false_func():
+            return layers.fill_constant(shape=[3, 2], dtype='int32', value=-1)
+
+        main_program = Program()
+        startup_program = Program()
+        with program_guard(main_program, startup_program):
+            x = layers.fill_constant(shape=[1], dtype='float32', value=0.1)
+            y = layers.fill_constant(shape=[1], dtype='float32', value=0.23)
+            pred = layers.less_than(y, x)
+            out = layers.cond(pred, true_func, false_func)
+            # out is one tensor
+
+        place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda(
+        ) else fluid.CPUPlace()
+        exe = fluid.Executor(place)
+        ret = exe.run(main_program, fetch_list=[out.name])
+        self.assertTrue(
+            np.allclose(np.asarray(ret), np.full((3, 2), -1, np.int32)))
+
+    def test_return_var_tuple(self):
+        """
+        pseudocode:
+
+        if True:
+            return 1, True
+        else:
+            return 3, 2
+        """
+
+        def true_func():
+            return layers.fill_constant(
+                shape=[1, 2], dtype='int32', value=1), layers.fill_constant(
+                    shape=[2, 3], dtype='bool', value=True)
+
+        def false_func():
+            return layers.fill_constant(
+                shape=[3, 4], dtype='float32', value=3), layers.fill_constant(
+                    shape=[4, 5], dtype='int64', value=2)
+
+        main_program = Program()
+        startup_program = Program()
+        with program_guard(main_program, startup_program):
+            pred = layers.fill_constant(shape=[1], dtype='bool', value=True)
+            out = layers.cond(pred, true_func, false_func)
+            # out is a tuple containing 2 tensors
+
+        place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda(
+        ) else fluid.CPUPlace()
+        exe = fluid.Executor(place)
+        ret = exe.run(main_program, fetch_list=out)
+        self.assertTrue(
+            np.allclose(np.asarray(ret[0]), np.full((1, 2), 1, np.int32)))
+        self.assertTrue(
+            np.allclose(np.asarray(ret[1]), np.full((2, 3), True, np.bool)))
+
+    def test_pass_and_modify_var(self):
+        """
+        pseudocode:
+        for i in range(5):
+            a = 7
+            if i % 2 == 0:
+                a = a * (i + 1)
+            else:
+                a = a - (i - 1)
+        """
+
+        def true_func(a, i):
+            a = a * (i + 1)
+            return a
+
+        def false_func(a, i):
+            a = a - (i - 1)
+            return a
+
+        main_program = Program()
+        startup_program = Program()
+        with program_guard(main_program, startup_program):
+            a = layers.fill_constant(shape=[3, 2, 1], dtype='int32', value=7)
+            i = fluid.data(name="i", shape=[1], dtype='int32')
+            pred = ((i % 2) == 0)
+            a = layers.cond(pred, lambda: true_func(a, i),
+                            lambda: false_func(a, i))
+        place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda(
+        ) else fluid.CPUPlace()
+        exe = fluid.Executor(place)
+        for feed_i in range(5):
+            expected_a = 7 * (feed_i + 1) if feed_i % 2 == 0 else 8 - feed_i
+            ret = exe.run(main_program,
+                          feed={'i': np.full((1), feed_i, np.int32)},
+                          fetch_list=[a])
+            self.assertTrue(
+                np.allclose(
+                    np.asarray(ret), np.full((3, 2, 1), expected_a, np.int32)))
+
+    def test_return_none(self):
+        """
+        pseudocode: test doing nothing in branches
+        for i in range(5):
+            if i % 2 == 0:
+                pass
+            else:
+                pass
+        """
+
+        def true_func():
+            pass
+
+        def false_func():
+            return None
+
+        main_program = Program()
+        startup_program = Program()
+        with program_guard(main_program, startup_program):
+            i = fluid.data(name="i", shape=[1], dtype='int32')
+            pred = ((i % 2) == 0)
+            out1 = layers.cond(pred, true_func, false_func)
+            out2 = layers.cond(pred, None, false_func)
+            out3 = layers.cond(pred, true_func, None)
+        place = fluid.CUDAPlace(0) if core.is_compiled_with_cuda(
+        ) else fluid.CPUPlace()
+        exe = fluid.Executor(place)
+        for feed_i in range(5):
+            # Test that output is None is runnable
+            exe.run(main_program, feed={'i': np.full((1), feed_i, np.int32)})
+            self.assertIsNone(out1)
+            self.assertIsNone(out2)
+            self.assertIsNone(out3)
+
+    def test_wrong_structure_exception(self):
+        """
+        test returning different number of tensors cannot merge into output
+        """
+
+        def func_return_none():
+            return None
+
+        def func_return_one_tensor():
+            return layers.fill_constant(shape=[2, 7], dtype='int32', value=3)
+
+        def func_return_two_tensors():
+            return layers.fill_constant(
+                shape=[3, 1], dtype='int32', value=7), layers.fill_constant(
+                    shape=[3, 1], dtype='int32', value=8)
+
+        main_program = Program()
+        startup_program = Program()
+        with program_guard(main_program, startup_program):
+            i = fluid.data(name="i", shape=[1], dtype='int32')
+            pred = ((i % 2) == 0)
+            with self.assertRaises(Exception) as e:
+                out = layers.cond(pred, i, func_return_one_tensor)
+            self.assertEqual("The true_fn in cond must be callable",
+                             str(e.exception))
+
+            with self.assertRaises(Exception) as e:
+                out = layers.cond(pred, func_return_one_tensor, np.asarray([3]))
+            self.assertEqual("The false_fn in cond must be callable",
+                             str(e.exception))
+
+            with self.assertRaises(Exception) as e:
+                out = layers.cond(pred, func_return_none,
+                                  func_return_one_tensor)
+            self.assertTrue(
+                "Incompatible return values of true_fn and false_fn in cond" in
+                str(e.exception))
+
+            with self.assertRaises(Exception) as e:
+                out = layers.cond(pred, func_return_two_tensors,
+                                  func_return_none)
+            self.assertTrue(
+                "Incompatible return values of true_fn and false_fn in cond" in
+                str(e.exception))
+
+            with self.assertRaises(Exception) as e:
+                out = layers.cond(pred, func_return_one_tensor,
+                                  func_return_two_tensors)
+            self.assertTrue(
+                "Incompatible return values of true_fn and false_fn in cond" in
+                str(e.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()