Tensor index (#16223)

* extend the slice function for python
test=develop

paddle/fluid/pybind/pybind.cc

@@ -347,7 +347,8 @@ PYBIND11_MODULE(core, m) {
       .def("_set_double_element", TensorSetElement<double>)
       .def("_get_double_element", TensorGetElement<double>)
       .def("_place", [](Tensor &self) { return self.place(); })
-      .def("_dtype", [](Tensor &self) { return self.type(); });
+      .def("_dtype", [](Tensor &self) { return self.type(); })
+      .def("__getitem__", PySliceTensor, py::return_value_policy::reference);
 
   py::class_<LoDTensor, Tensor>(m, "LoDTensor", R"DOC(
     LoDTensor is a Tensor with optional LoD information.
@@ -499,6 +500,13 @@ PYBIND11_MODULE(core, m) {
 
            Returns:
                out (bool): whether the lod is valid.
+           )DOC")
+      .def("__getitem__", PySliceTensor, py::return_value_policy::reference,
+           R"DOC(
+           Slice the original Tensor, and remove the LoD information.
+
+           Returns:
+               out (Tensor): new Tensor(NOT LoDTensor).
            )DOC");
 
   py::class_<SelectedRows>(m, "SelectedRows")

python/paddle/fluid/framework.py

@@ -627,6 +627,183 @@ class Variable(object):
         """
         self.error_clip = error_clip
 
+    def _slice_indices(self, slice, length):
+        """
+        Reference implementation for the slice.indices method.
+        """
+        # Compute step and length as integers.
+        step = 1 if slice.step is None else slice.step
+
+        # Raise ValueError for negative length or zero step.
+        if length < 0:
+            raise ValueError("length should not be negative")
+        if step == 0:
+            raise ValueError("slice step cannot be zero")
+
+        # Find lower and upper bounds for start and stop.
+        lower = -1 if step < 0 else 0
+        upper = length - 1 if step < 0 else length
+
+        # Compute start.
+        if slice.start is None:
+            start = upper if step < 0 else lower
+        else:
+            start = slice.start
+            start = max(start + length, lower) if start < 0 else min(start,
+                                                                     upper)
+
+        # Compute stop.
+        if slice.stop is None:
+            stop = lower if step < 0 else upper
+        else:
+            stop = slice.stop
+            stop = max(stop + length, lower) if stop < 0 else min(stop, upper)
+
+        return start, stop, step
+
+    def _detectEllipsis(self, item):
+        has_ellipsis = False
+        start = 0
+        end = len(self.shape)
+        for index, o in enumerate(item):
+            if o is Ellipsis:
+                if has_ellipsis:
+                    raise ValueError("Index can have one ellipsis only.")
+                has_ellipsis = True
+                start = index
+            else:
+                if has_ellipsis:
+                    end = index
+        return has_ellipsis, start, end
+
+    def _reconstructSliceinfo(self, item):
+        has_ellipsis, start, end = self._detectEllipsis(item)
+        if has_ellipsis:
+            newitem = []
+            for i in range(start):
+                newitem.append(item[i])
+            for i in range(start, end):
+                newitem.append(slice(None, None, None))
+            for i in range(end, len(item)):
+                newitem.append(item[i])
+            return newitem
+        else:
+            return None
+
+    def _detectContinuesSlice(self, item):
+        starts = []
+        ends = []
+        for index, o in enumerate(item):
+            if isinstance(o, int):
+                start = int(o)
+                if (index > 0 and index >= self.shape[index]) \
+                        or (index < 0 and (index + self.shape[index]) < 0):
+                    raise IndexError("invalid index")
+                start = max(start + self.shape[index], 0) if start < 0 else min(
+                    start, self.shape[index])
+                starts.append(start)
+                ends.append(start + 1)
+            elif isinstance(o, slice):
+                start, stop, step = self._slice_indices(o, self.shape[index])
+                if step == 1 or step == -1:
+                    starts.append(start)
+                    ends.append(stop)
+                else:
+                    return False, None
+            else:
+                raise IndexError("Valid index accept int or slice or ellipsis")
+        return True, [starts, ends]
+
+    def _cloneVar(self, copy=False):
+        if not copy:
+            return self.block.create_var(
+                name=unique_name.generate(".".join(self.name)),
+                dtype=self.dtype,
+                persistable=self.persistable,
+                stop_gradient=self._stop_gradient, )
+        else:
+            return self
+
+    def _sliceVar(self, axes, starts, ends):
+        new_var = self._cloneVar()
+        self.block.append_op(
+            type="slice",
+            inputs={'Input': [self]},
+            outputs={'Out': [new_var]},
+            attrs={'axes': axes,
+                   'starts': starts,
+                   'ends': ends})
+        return new_var
+
+    def _concatVar(self, inputs, axis):
+        new_var = self._cloneVar()
+        self.block.append_op(
+            type="concat",
+            inputs={'X': inputs},
+            outputs={'Out': [new_var]},
+            attrs={'axis': axis, })
+        return new_var
+
+    def _sliceAndConcatVar(self, item, axis):
+        if isinstance(item, slice):
+            if self.shape[axis] < 0:
+                return self._cloneVar(True)
+            start, stop, step = self._slice_indices(item, self.shape[axis])
+            if step == 1:
+                return self._sliceVar([axis], [start], [stop])
+            else:
+                vars = []
+                if step > 0:
+                    while start < stop:
+                        vars.append(
+                            self._sliceVar([axis], [start], [start + 1]))
+                        start += step
+                else:
+                    while start > stop:
+                        vars.append(
+                            self._sliceVar([axis], [start], [start + 1]))
+                        start += step
+                return self._concatVar(vars, axis)
+        elif isinstance(item, int):
+            if self.shape[axis] < 0:
+                return self._cloneVar(True)
+            index = int(item)
+            if (index > 0 and index >= self.shape[axis])\
+                    or (index < 0 and (index + self.shape[axis]) < 0):
+                raise IndexError("invalid index")
+            return self._sliceVar([axis], [index], [index + 1])
+        else:
+            raise IndexError("Valid index accept int or slice or tuple")
+
+    def __getitem__(self, item):
+        """
+        Slice the variable.
+
+        Args:
+            item(int/slice/tuple) : the index.
+
+        Returns:
+            Sliced variable
+        """
+        new_var = None
+        if isinstance(item, tuple):
+            if len(item) > len(self.shape):
+                raise IndexError("Too many indexes")
+            newitem = self._reconstructSliceinfo(item) or item
+            check, info = self._detectContinuesSlice(newitem)
+            if check:
+                starts = info[0]
+                ends = info[1]
+                axes = [i for i in range(len(starts))]
+                return self._sliceVar(axes, starts, ends)
+            else:
+                new_var = self
+                for index, o in enumerate(newitem):
+                    new_var = new_var._sliceAndConcatVar(o, index)
+        else:
+            new_var = self._sliceAndConcatVar(item, 0)
+        return new_var
+
 
 def get_all_op_protos():
     """

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

@@ -14,6 +14,7 @@
 
 from __future__ import print_function
 
+import paddle.fluid as fluid
 import paddle.fluid.core as core
 import unittest
 import numpy
@@ -183,6 +184,58 @@ class TestTensor(unittest.TestCase):
             tensor_array = numpy.array(tensor)
             self.assertEqual((0, 1), tensor_array.shape)
 
+    def run_sliece_tensor(self, place):
+
+        tensor = fluid.Tensor()
+        shape = [3, 3, 3]
+        tensor._set_dims(shape)
+
+        tensor_array = numpy.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                                    [[10, 11, 12], [13, 14, 15], [16, 17, 18]],
+                                    [[19, 20, 21], [22, 23, 24], [25, 26, 27]]])
+
+        tensor.set(tensor_array, place)
+        n1 = tensor[1]
+        t1 = tensor_array[1]
+        self.assertTrue((numpy.array(n1) == numpy.array(t1)).all())
+
+        n2 = tensor[1:]
+        t2 = tensor_array[1:]
+        self.assertTrue((numpy.array(n2) == numpy.array(t2)).all())
+
+        n3 = tensor[0:2:]
+        t3 = tensor_array[0:2:]
+        self.assertTrue((numpy.array(n3) == numpy.array(t3)).all())
+
+        n4 = tensor[2::-2]
+        t4 = tensor_array[2::-2]
+        self.assertTrue((numpy.array(n4) == numpy.array(t4)).all())
+
+        n5 = tensor[2::-2][0]
+        t5 = tensor_array[2::-2][0]
+        self.assertTrue((numpy.array(n5) == numpy.array(t5)).all())
+
+        n6 = tensor[2:-1:-1]
+        t6 = tensor_array[2:-1:-1]
+        self.assertTrue((numpy.array(n6) == numpy.array(t6)).all())
+
+        n7 = tensor[0:, 0:]
+        t7 = tensor_array[0:, 0:]
+        self.assertTrue((numpy.array(n7) == numpy.array(t7)).all())
+
+        n8 = tensor[0::1, 0::-1, 2:]
+        t8 = tensor_array[0::1, 0::-1, 2:]
+        self.assertTrue((numpy.array(n8) == numpy.array(t8)).all())
+
+    def test_sliece_tensor(self):
+        # run cpu first
+        place = core.CPUPlace()
+        self.run_sliece_tensor(place)
+
+        if core.is_compiled_with_cuda():
+            place = core.CUDAPlace(0)
+            self.run_sliece_tensor(place)
+
 
 if __name__ == '__main__':
     unittest.main()

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

@@ -16,8 +16,10 @@ from __future__ import print_function
 
 import unittest
 from paddle.fluid.framework import default_main_program, Program, convert_np_dtype_to_dtype_
+import paddle.fluid as fluid
 import paddle.fluid.core as core
 import numpy as np
+from test_imperative_base import new_program_scope
 
 
 class TestVariable(unittest.TestCase):
@@ -60,6 +62,100 @@ class TestVariable(unittest.TestCase):
             name='step_scopes', type=core.VarDesc.VarType.STEP_SCOPES)
         self.assertEqual(core.VarDesc.VarType.STEP_SCOPES, var.type)
 
+    def _test_slice(self):
+        b = default_main_program().current_block()
+        w = b.create_var(dtype="float64", shape=[784, 100, 100], lod_level=0)
+
+        for i in range(3):
+            nw = w[i]
+            self.assertEqual((1, 100, 100), nw.shape)
+
+        nw = w[:]
+        self.assertEqual((784, 100, 100), nw.shape)
+
+        nw = w[:, :, ...]
+        self.assertEqual((784, 100, 100), nw.shape)
+
+        nw = w[::2, ::2, :]
+        self.assertEqual((392, 50, 100), nw.shape)
+
+        nw = w[::-2, ::-2, :]
+        self.assertEqual((392, 50, 100), nw.shape)
+
+        self.assertEqual(0, nw.lod_level)
+
+        place = fluid.CPUPlace()
+        main = fluid.Program()
+        with fluid.program_guard(main):
+            exe = fluid.Executor(place)
+            tensor_array = np.array(
+                [[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                 [[10, 11, 12], [13, 14, 15], [16, 17, 18]],
+                 [[19, 20, 21], [22, 23, 24], [25, 26, 27]]]).astype('float32')
+            var = fluid.layers.assign(tensor_array)
+            var1 = var[0, 1, 1]
+            var2 = var[1:]
+            var3 = var[0:1]
+            var4 = var[..., ]
+            var5 = var[2::-2]
+            var6 = var[1, 1:, 1:]
+            var7 = var[1, ..., 1:]
+            var8 = var[1, ...]
+            local_out = exe.run(main,
+                                fetch_list=[
+                                    var, var1, var2, var3, var4, var5, var6,
+                                    var7, var8
+                                ])
+
+            self.assertTrue((np.array(local_out[1]) == np.array(tensor_array[
+                0, 1, 1])).all())
+            self.assertTrue((np.array(local_out[2]) == np.array(tensor_array[
+                1:])).all())
+            self.assertTrue((np.array(local_out[3]) == np.array(tensor_array[
+                0:1])).all())
+            self.assertTrue((np.array(local_out[4]) == np.array(
+                tensor_array[..., ])).all())
+            self.assertTrue((np.array(local_out[5]) == np.array(tensor_array[
+                2::-2])).all())
+            self.assertTrue((np.array(local_out[6]) == np.array(tensor_array[
+                1, 1:, 1:])).all())
+            self.assertTrue((np.array(local_out[7]) == np.array(tensor_array[
+                1, ..., 1:])).all())
+            self.assertTrue((np.array(local_out[8]) == np.array(tensor_array[
+                1, ...])).all())
+
+    def test_slice(self):
+        self._test_slice()
+
+
+class TestVariableImperative(unittest.TestCase):
+    def _test_slice(self):
+        b = default_main_program().current_block()
+        w = b.create_var(dtype="float64", shape=[784, 100, 100], lod_level=0)
+
+        for i in range(3):
+            nw = w[i]
+            self.assertEqual([1, 100, 100], nw.shape)
+
+        nw = w[:]
+        self.assertEqual([784, 100, 100], nw.shape)
+
+        nw = w[:, :, :]
+        self.assertEqual([784, 100, 100], nw.shape)
+
+        nw = w[::2, ::2, :]
+        self.assertEqual([392, 50, 100], nw.shape)
+
+        nw = w[::-2, ::-2, :]
+        self.assertEqual([392, 50, 100], nw.shape)
+
+        nw = w[0::-2, 0::-2, :]
+        self.assertEqual([1, 1, 100], nw.shape)
+
+    def test_slice(self):
+        with fluid.imperative.guard():
+            self._test_slice()
+
 
 if __name__ == '__main__':
     unittest.main()