Reverse Sequence (#920)

* sequence reverse

* fix test

* more tests for reverse_sequence

* remove debug prints, change perms

* fix formatting; remove dead code

* make seq_lengths a parameter

* autodiff + tests

src/ngraph/CMakeLists.txt

@@ -84,6 +84,7 @@ set (SRC
     op/reshape.cpp
     op/result.cpp
     op/reverse.cpp
+    op/reverse_sequence.cpp
     op/select_and_scatter.cpp
     op/select.cpp
     op/sign.cpp

src/ngraph/op/reverse_sequence.cpp

+/*******************************************************************************
+* Copyright 2017-2018 Intel Corporation
+*
+* 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 <algorithm>
+#include <memory>
+#include <typeindex>
+#include <typeinfo>
+
+#include "ngraph/node.hpp"
+#include "ngraph/op/reverse_sequence.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+op::ReverseSequence::ReverseSequence(const std::shared_ptr<Node> arg,
+                                     const std::shared_ptr<Node> seq_indices,
+                                     size_t batch_axis,
+                                     size_t seq_axis)
+    : RequiresTensorViewArgs("ReverseSequence", {arg, seq_indices})
+    , m_batch_axis(batch_axis)
+    , m_seq_axis(seq_axis)
+{
+    if (seq_indices->get_shape().size() != 1)
+    {
+        throw ngraph_error("indices should be a 1-dimensional array");
+    }
+
+    if (arg->get_shape().at(batch_axis) != seq_indices->get_shape().at(0))
+    {
+        throw ngraph_error("Sequence length size should be equal to batch axis dimension");
+    }
+
+    set_value_type_checked(arg->get_element_type(), arg->get_shape());
+}
+
+shared_ptr<Node> op::ReverseSequence::copy_with_new_args(const NodeVector& new_args) const
+{
+    if (new_args.size() != 2)
+    {
+        throw ngraph_error("Incorrect number of new arguments");
+    }
+
+    auto res =
+        make_shared<ReverseSequence>(new_args.at(0), new_args.at(1), m_batch_axis, m_seq_axis);
+    return res;
+}
+
+void op::ReverseSequence::generate_adjoints(autodiff::Adjoints& adjoints, const NodeVector& deltas)
+{
+    auto x = get_argument(0);
+    auto rs_delta =
+        make_shared<ReverseSequence>(deltas.at(0), get_argument(1), m_batch_axis, m_seq_axis);
+    adjoints.add_delta(x, rs_delta);
+}

src/ngraph/op/reverse_sequence.hpp

+/*******************************************************************************
+* Copyright 2017-2018 Intel Corporation
+*
+* 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.
+*******************************************************************************/
+
+#pragma once
+
+#include <memory>
+
+#include "ngraph/op/util/requires_tensor_view_args.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        class ReverseSequence : public util::RequiresTensorViewArgs
+        {
+        public:
+            /// \brief Constructs an arcsin operation.
+            ///
+            /// \param arg Node that produces the input tensor.
+            ReverseSequence(const std::shared_ptr<Node> arg,
+                            const std::shared_ptr<Node> seq_lengths,
+                            size_t batch_axis,
+                            size_t seq_axis);
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override;
+
+            size_t get_batch_axis() const { return m_batch_axis; }
+            size_t get_sequence_axis() const { return m_seq_axis; }
+        protected:
+            virtual void generate_adjoints(autodiff::Adjoints& adjoints,
+                                           const NodeVector& deltas) override;
+
+        private:
+            size_t m_batch_axis{0};
+            size_t m_seq_axis{0};
+        };
+    }
+}

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -76,6 +76,7 @@
 #include "ngraph/op/reshape.hpp"
 #include "ngraph/op/result.hpp"
 #include "ngraph/op/reverse.hpp"
+#include "ngraph/op/reverse_sequence.hpp"
 #include "ngraph/op/select.hpp"
 #include "ngraph/op/select_and_scatter.hpp"
 #include "ngraph/op/sign.hpp"
@@ -2783,6 +2784,21 @@ namespace ngraph
                 writer << "                {" << join(reverse->get_reversed_axes()) << "});\n";
             }
 
+            template <>
+            void CPU_Emitter::EMITTER_DECL(ngraph::op::ReverseSequence)
+            {
+                auto rs = static_cast<const ngraph::op::ReverseSequence*>(node);
+                auto arg_shape = args[0].get_shape();
+
+                writer << "reference::reverse_sequence<" << out[0].get_type() << ","
+                       << args[1].get_type() << ">(" << args[0].get_name() << ",\n";
+                writer << "                " << out[0].get_name() << ",\n";
+                writer << "                {" << join(arg_shape) << "},\n";
+                writer << "                 " << rs->get_batch_axis() << ",\n";
+                writer << "                 " << rs->get_sequence_axis() << ",\n";
+                writer << "                 " << args[1].get_name() << ");\n";
+            }
+
             template <>
             void CPU_Emitter::EMITTER_DECL(ngraph::op::ReduceWindow)
             {

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -86,6 +86,7 @@
 #include "ngraph/op/reshape.hpp"
 #include "ngraph/op/result.hpp"
 #include "ngraph/op/reverse.hpp"
+#include "ngraph/op/reverse_sequence.hpp"
 #include "ngraph/op/select.hpp"
 #include "ngraph/op/select_and_scatter.hpp"
 #include "ngraph/op/sign.hpp"
@@ -262,6 +263,7 @@ static const runtime::cpu::OpMap dispatcher{
     {TI(ngraph::op::Not), &runtime::cpu::CPU_Emitter::emit<op::Not>},
     {TI(ngraph::op::MaxPool), &runtime::cpu::CPU_Emitter::emit<op::MaxPool>},
     {TI(ngraph::op::Reverse), &runtime::cpu::CPU_Emitter::emit<op::Reverse>},
+    {TI(ngraph::op::ReverseSequence), &runtime::cpu::CPU_Emitter::emit<op::ReverseSequence>},
     {TI(ngraph::op::Result), &runtime::cpu::CPU_Emitter::emit<op::Result>},
     {TI(ngraph::op::ReduceWindow), &runtime::cpu::CPU_Emitter::emit<op::ReduceWindow>},
     {TI(ngraph::op::SelectAndScatter), &runtime::cpu::CPU_Emitter::emit<op::SelectAndScatter>},
@@ -363,6 +365,7 @@ void runtime::cpu::CPU_ExternalFunction::compile()
 #include "ngraph/runtime/reference/reshape.hpp"
 #include "ngraph/runtime/reference/result.hpp"
 #include "ngraph/runtime/reference/reverse.hpp"
+#include "ngraph/runtime/reference/reverse_sequence.hpp"
 #include "ngraph/runtime/reference/select_and_scatter.hpp"
 #include "ngraph/runtime/reference/slice.hpp"
 #include "ngraph/runtime/reference/sum.hpp"

src/ngraph/runtime/reference/reverse_sequence.hpp

+/*******************************************************************************
+* Copyright 2017-2018 Intel Corporation
+*
+* 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.
+*******************************************************************************/
+
+#pragma once
+
+#include <cmath>
+#include <numeric>
+
+#include "ngraph/coordinate_transform.hpp"
+#include "ngraph/util.hpp"
+
+namespace ngraph
+{
+    namespace runtime
+    {
+        namespace reference
+        {
+            template <typename T, typename U>
+            void reverse_sequence(const T* arg,
+                                  T* out,
+                                  const Shape& arg_shape,
+                                  size_t batch_axis,
+                                  size_t sequence_axis,
+                                  U* sequence_lengths)
+            {
+                CoordinateTransform input_transform(arg_shape);
+                for (const Coordinate& in_coord : input_transform)
+                {
+                    size_t batch_index = in_coord[batch_axis];
+                    auto orig_seq_index = static_cast<size_t>(sequence_lengths[batch_index]);
+
+                    if (orig_seq_index > arg_shape.at(sequence_axis))
+                    {
+                        throw ngraph_error(
+                            "One of the elements of sequence lengths is greater than sequence axis "
+                            "dimension");
+                    }
+
+                    if (orig_seq_index == 0)
+                    {
+                        orig_seq_index = 1;
+                    }
+
+                    size_t sequence_index = in_coord[sequence_axis] < orig_seq_index
+                                                ? orig_seq_index - in_coord[sequence_axis] - 1
+                                                : in_coord[sequence_axis];
+
+                    //make a copy of in_coord and update sequence_index
+                    Coordinate out_coord = in_coord;
+                    out_coord[sequence_axis] = sequence_index;
+                    out[input_transform.index(out_coord)] = arg[input_transform.index(in_coord)];
+                }
+            }
+        }
+    }
+}

test/cpu_test.cpp

@@ -29,6 +29,7 @@
 #include "ngraph/op/batch_norm.hpp"
 #include "ngraph/op/get_output_element.hpp"
 #include "ngraph/op/parameter.hpp"
+#include "ngraph/op/reverse_sequence.hpp"
 #include "ngraph/pass/manager.hpp"
 #include "ngraph/pass/visualize_tree.hpp"
 #include "ngraph/serializer.hpp"
@@ -60,3 +61,216 @@ TEST(cpu_test, unhandled_op)
     auto backend = runtime::Backend::create("CPU");
     ASSERT_THROW(backend->compile(f), ngraph_error);
 }
+
+TEST(cpu_test, reverse_sequence_n2c3h4w2)
+{
+    Shape shape{2, 3, 4, 2};
+    Shape seq_len_shape{4};
+    auto A = make_shared<op::Parameter>(element::i32, shape);
+    auto B = make_shared<op::Parameter>(element::i32, seq_len_shape);
+
+    size_t batch_axis = 2;
+    size_t sequence_axis = 1;
+    auto rs = std::make_shared<op::ReverseSequence>(A, B, batch_axis, sequence_axis);
+
+    auto f = make_shared<Function>(rs, op::ParameterVector{A, B});
+
+    auto backend = runtime::Backend::create("CPU");
+
+    // Create some tensors for input/output
+    shared_ptr<runtime::TensorView> a = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> b = backend->create_tensor(element::i32, seq_len_shape);
+
+    shared_ptr<runtime::TensorView> result = backend->create_tensor(element::i32, shape);
+
+    std::vector<int> input{
+        0,  0, 3,  0, 6,  0, 9,  0, 1,  0, 4,  0, 7,  0, 10, 0, 2,  0, 5,  0, 8,  0, 11, 0,
+        12, 0, 15, 0, 18, 0, 21, 0, 13, 0, 16, 0, 19, 0, 22, 0, 14, 0, 17, 0, 20, 0, 23, 0,
+    };
+
+    std::vector<int> seq_lenghts{1, 2, 1, 2};
+    copy_data(b, seq_lenghts);
+
+    std::vector<int> expected{
+        0,  0, 4,  0, 6,  0, 10, 0, 1,  0, 3,  0, 7,  0, 9,  0, 2,  0, 5,  0, 8,  0, 11, 0,
+
+        12, 0, 16, 0, 18, 0, 22, 0, 13, 0, 15, 0, 19, 0, 21, 0, 14, 0, 17, 0, 20, 0, 23, 0};
+
+    copy_data(a, input);
+
+    backend->call(f, {result}, {a, b});
+    EXPECT_EQ(read_vector<int>(result), expected);
+}
+
+TEST(cpu_test, reverse_sequence_n4c3h2w2)
+{
+    Shape shape{4, 3, 2, 2};
+    auto A = make_shared<op::Parameter>(element::i32, shape);
+    Shape seq_len_shape{4};
+    auto B = make_shared<op::Parameter>(element::i32, seq_len_shape);
+
+    size_t batch_axis = 0;
+    size_t sequence_axis = 1;
+
+    auto rs = std::make_shared<op::ReverseSequence>(A, B, batch_axis, sequence_axis);
+
+    auto f = make_shared<Function>(rs, op::ParameterVector{A, B});
+
+    auto backend = runtime::Backend::create("CPU");
+
+    // Create some tensors for input/output
+    shared_ptr<runtime::TensorView> a = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> b = backend->create_tensor(element::i32, seq_len_shape);
+
+    shared_ptr<runtime::TensorView> result = backend->create_tensor(element::i32, shape);
+
+    std::vector<int> seq_lenghts{1, 2, 3, 3};
+    copy_data(b, seq_lenghts);
+
+    std::vector<int> input{0,  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, 28, 29, 30, 31,
+                           32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47};
+
+    std::vector<int> expected{0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 16, 17, 18, 19,
+                              12, 13, 14, 15, 20, 21, 22, 23, 32, 33, 34, 35, 28, 29, 30, 31,
+                              24, 25, 26, 27, 44, 45, 46, 47, 40, 41, 42, 43, 36, 37, 38, 39};
+
+    copy_data(a, input);
+
+    backend->call(f, {result}, {a, b});
+    EXPECT_EQ(read_vector<int>(result), expected);
+}
+
+TEST(cpu_test, reverse_sequence_n4d2c3h2w2)
+{
+    Shape shape{4, 2, 3, 2, 2};
+    auto A = make_shared<op::Parameter>(element::i32, shape);
+    Shape seq_len_shape{4};
+    auto B = make_shared<op::Parameter>(element::i32, seq_len_shape);
+
+    size_t batch_axis = 0;
+    size_t sequence_axis = 2;
+
+    auto rs = std::make_shared<op::ReverseSequence>(A, B, batch_axis, sequence_axis);
+
+    auto f = make_shared<Function>(rs, op::ParameterVector{A, B});
+
+    auto backend = runtime::Backend::create("CPU");
+
+    // Create some tensors for input/output
+    shared_ptr<runtime::TensorView> a = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> b = backend->create_tensor(element::i32, seq_len_shape);
+
+    shared_ptr<runtime::TensorView> result = backend->create_tensor(element::i32, shape);
+
+    std::vector<int> input{0,  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, 28, 29, 30, 31,
+                           32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
+                           48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
+                           64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79,
+                           80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95};
+
+    std::vector<int> expected{0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15,
+                              16, 17, 18, 19, 20, 21, 22, 23, 28, 29, 30, 31, 24, 25, 26, 27,
+                              32, 33, 34, 35, 40, 41, 42, 43, 36, 37, 38, 39, 44, 45, 46, 47,
+                              48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
+                              64, 65, 66, 67, 68, 69, 70, 71, 76, 77, 78, 79, 72, 73, 74, 75,
+                              80, 81, 82, 83, 88, 89, 90, 91, 84, 85, 86, 87, 92, 93, 94, 95};
+
+    copy_data(a, input);
+
+    std::vector<int> seq_lenghts{1, 2, 1, 2};
+    copy_data(b, seq_lenghts);
+
+    backend->call(f, {result}, {a, b});
+    EXPECT_EQ(read_vector<int>(result), expected);
+}
+
+TEST(cpu_test, backwards_reverse_sequence_n3_c2_h3)
+{
+    Shape shape{3, 2, 3};
+    auto backend = runtime::Backend::create("CPU");
+
+    auto A = make_shared<op::Parameter>(element::i32, shape);
+    Shape seq_len_shape{2};
+    auto B = make_shared<op::Parameter>(element::i32, seq_len_shape);
+
+    size_t batch_axis = 1;
+    size_t sequence_axis = 0;
+
+    auto rs = std::make_shared<op::ReverseSequence>(A, B, batch_axis, sequence_axis);
+    auto f = make_shared<Function>(rs, op::ParameterVector{A, B});
+
+    shared_ptr<runtime::TensorView> a = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> b = backend->create_tensor(element::i32, seq_len_shape);
+    shared_ptr<runtime::TensorView> c = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> da = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> db = backend->create_tensor(element::i32, seq_len_shape);
+
+    //input values don't matter
+    vector<int> va(shape_size(shape), 0);
+
+    vector<int> vc{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18};
+    vector<int> expected{13, 14, 15, 16, 17, 18, 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6};
+
+    copy_data(c, vc);
+    copy_data(a, va);
+
+    std::vector<int> seq_lenghts{3, 3};
+    copy_data(b, seq_lenghts);
+
+    auto C = make_shared<op::Parameter>(element::i32, shape);
+    auto df = autodiff::backprop_function(f);
+    backend->call(df, {da, db}, {a, b, c});
+    ASSERT_EQ(read_vector<int>(da), expected);
+}
+
+TEST(cpu_test, backwards_reverse_sequence_n4d2c3h2w2)
+{
+    Shape shape{4, 2, 3, 2, 2};
+    auto backend = runtime::Backend::create("CPU");
+
+    auto A = make_shared<op::Parameter>(element::i32, shape);
+    Shape seq_len_shape{4};
+    auto B = make_shared<op::Parameter>(element::i32, seq_len_shape);
+
+    size_t batch_axis = 0;
+    size_t sequence_axis = 2;
+
+    auto rs = std::make_shared<op::ReverseSequence>(A, B, batch_axis, sequence_axis);
+    auto f = make_shared<Function>(rs, op::ParameterVector{A, B});
+
+    shared_ptr<runtime::TensorView> a = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> b = backend->create_tensor(element::i32, seq_len_shape);
+    shared_ptr<runtime::TensorView> c = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> da = backend->create_tensor(element::i32, shape);
+    shared_ptr<runtime::TensorView> db = backend->create_tensor(element::i32, seq_len_shape);
+
+    //input values don't matter
+    vector<int> va(shape_size(shape), 0);
+
+    std::vector<int> vc{0,  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, 28, 29, 30, 31,
+                        32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
+                        48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
+                        64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79,
+                        80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95};
+
+    std::vector<int> expected{0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15,
+                              16, 17, 18, 19, 20, 21, 22, 23, 28, 29, 30, 31, 24, 25, 26, 27,
+                              32, 33, 34, 35, 40, 41, 42, 43, 36, 37, 38, 39, 44, 45, 46, 47,
+                              48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
+                              64, 65, 66, 67, 68, 69, 70, 71, 76, 77, 78, 79, 72, 73, 74, 75,
+                              80, 81, 82, 83, 88, 89, 90, 91, 84, 85, 86, 87, 92, 93, 94, 95};
+
+    copy_data(c, vc);
+    copy_data(a, va);
+
+    std::vector<int> seq_lenghts{1, 2, 1, 2};
+    copy_data(b, seq_lenghts);
+
+    auto C = make_shared<op::Parameter>(element::i32, shape);
+    auto df = autodiff::backprop_function(f);
+    backend->call(df, {da, db}, {a, b, c});
+    ASSERT_EQ(read_vector<int>(da), expected);
+}