Merge pull request #12884 from sneaxiy/sequence_mask_op

Add sequence_mask_op for DAM model

paddle/fluid/API.spec

@@ -162,6 +162,7 @@ paddle.fluid.layers.crop ArgSpec(args=['x', 'shape', 'offsets', 'name'], varargs
 paddle.fluid.layers.rank_loss ArgSpec(args=['label', 'left', 'right', 'name'], varargs=None, keywords=None, defaults=(None,))
 paddle.fluid.layers.prelu ArgSpec(args=['x', 'mode', 'param_attr', 'name'], varargs=None, keywords=None, defaults=(None, None))
 paddle.fluid.layers.flatten ArgSpec(args=['x', 'axis', 'name'], varargs=None, keywords=None, defaults=(1, None))
+paddle.fluid.layers.sequence_mask ArgSpec(args=['x', 'maxlen', 'dtype', 'name'], varargs=None, keywords=None, defaults=(None, 'int64', None))
 paddle.fluid.layers.stack ArgSpec(args=['x', 'axis'], varargs=None, keywords=None, defaults=(0,))
 paddle.fluid.layers.data ArgSpec(args=['name', 'shape', 'append_batch_size', 'dtype', 'lod_level', 'type', 'stop_gradient'], varargs=None, keywords=None, defaults=(True, 'float32', 0, VarType.LOD_TENSOR, True))
 paddle.fluid.layers.open_recordio_file ArgSpec(args=['filename', 'shapes', 'lod_levels', 'dtypes', 'pass_num', 'for_parallel'], varargs=None, keywords=None, defaults=(1, True))

paddle/fluid/operators/sequence_mask_op.cc

@@ -0,0 +1,26 @@
+// 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.
+
+#include "paddle/fluid/operators/sequence_mask_op.h"
+
+REGISTER_OPERATOR(sequence_mask, paddle::operators::SequenceMaskOp,
+                  paddle::operators::SequenceMaskOpMaker,
+                  paddle::framework::EmptyGradOpMaker);
+
+REGISTER_OP_CPU_KERNEL(
+    sequence_mask,
+    paddle::operators::SequenceMaskKernel<paddle::platform::CPUDeviceContext,
+                                          int>,
+    paddle::operators::SequenceMaskKernel<paddle::platform::CPUDeviceContext,
+                                          int64_t>);

paddle/fluid/operators/sequence_mask_op.cu

@@ -0,0 +1,22 @@
+// 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.
+
+#include "paddle/fluid/operators/sequence_mask_op.h"
+
+REGISTER_OP_CUDA_KERNEL(
+    sequence_mask,
+    paddle::operators::SequenceMaskKernel<paddle::platform::CUDADeviceContext,
+                                          int>,
+    paddle::operators::SequenceMaskKernel<paddle::platform::CUDADeviceContext,
+                                          int64_t>);

paddle/fluid/operators/sequence_mask_op.h

@@ -0,0 +1,154 @@
+// 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.
+
+#pragma once
+
+#ifdef __NVCC__
+#include <thrust/device_ptr.h>
+#include <thrust/functional.h>
+#include <thrust/reduce.h>
+#else
+#include <algorithm>
+#endif
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/platform/for_range.h"
+
+namespace paddle {
+namespace operators {
+
+class SequenceMaskOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext *ctx) const override {
+    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) must exist");
+    PADDLE_ENFORCE(ctx->HasOutput("Y"), "Output(Y) must exist");
+
+    auto maxlen = ctx->Attrs().Get<int>("maxlen");
+    if (maxlen > 0) {  // We can only infershape when maxlen > 0
+      auto dim = framework::vectorize2int(ctx->GetInputDim("X"));
+      dim.push_back(maxlen);
+      ctx->SetOutputDim("Y", framework::make_ddim(dim));
+    }
+  }
+};
+
+class SequenceMaskOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X", "The input tensor of sequence_mask op.");
+    AddOutput("Y", "The output mask of sequence_mask op.");
+    AddAttr<int>("maxlen",
+                 "The maximum length of the sequence. If maxlen < 0, maxlen "
+                 "= max(Input(X)).")
+        .SetDefault(-1)
+        .AddCustomChecker([](int &v) {
+          PADDLE_ENFORCE(v < 0 || v >= 1,
+                         "Attr(maxlen) must be less than 0 or larger than 1");
+        });
+    AddAttr<int>("out_dtype", "Output data type");
+    AddComment(R"DOC(
+SequenceMask Operator
+
+This operator outputs a Mask according to Input(X) and Attr(maxlen).
+Supposing Input(X) is a Tensor with shape [d_1, d_2, ..., d_n], the
+Output(Y) is a mask with shape [d_1, d_2, ..., d_n, maxlen], where:
+
+Y(i_1, i_2, ..., i_n, j) = (j < X(i_1, i_2, ..., i_n)) 
+
+If maxlen < 0, maxlen = max(X)
+    )DOC");
+  }
+};
+
+template <typename Tx, typename Ty>
+struct SequenceMaskForRangeFunctor {
+  HOSTDEVICE SequenceMaskForRangeFunctor(const Tx *x, Ty *y, int maxlen)
+      : x_(x), y_(y), maxlen_(maxlen) {}
+
+  HOSTDEVICE void operator()(int y_idx) const {
+    int x_idx = y_idx / maxlen_;
+    int j = y_idx % maxlen_;
+    y_[y_idx] = static_cast<Ty>(j < x_[x_idx] ? 1 : 0);
+  }
+
+ private:
+  const Tx *x_;
+  Ty *y_;
+  int maxlen_;
+};
+
+template <typename DeviceContext, typename Tx>
+struct SequenceMaskFunctor {
+  using Tensor = framework::LoDTensor;
+
+  SequenceMaskFunctor(const DeviceContext &ctx, const Tx *x, Tensor *y,
+                      int limits, int maxlen)
+      : ctx_(ctx), x_(x), y_(y), limits_(limits), maxlen_(maxlen) {}
+
+  template <typename Ty>
+  void operator()() const {
+    auto *y_data = y_->mutable_data<Ty>(ctx_.GetPlace());
+    platform::ForRange<DeviceContext> for_range(ctx_, limits_);
+    for_range(SequenceMaskForRangeFunctor<Tx, Ty>(x_, y_data, maxlen_));
+  }
+
+ private:
+  const DeviceContext &ctx_;
+  const Tx *x_;
+  Tensor *y_;
+  int limits_;
+  int maxlen_;
+};
+
+template <typename DeviceContext, typename Tx>
+class SequenceMaskKernel : public framework::OpKernel<Tx> {
+  using Tensor = framework::LoDTensor;
+
+ public:
+  void Compute(const framework::ExecutionContext &ctx) const override {
+    auto *x = ctx.Input<Tensor>("X");
+    auto *y = ctx.Output<Tensor>("Y");
+    auto maxlen = ctx.Attr<int>("maxlen");
+
+    auto *x_data = x->data<Tx>();
+    auto x_numel = x->numel();
+    if (maxlen < 0) {
+#ifdef __NVCC__
+      VLOG(10)
+          << "SequenceMaskOp on GPU may be slow when maxlen is not provided.";
+      maxlen = static_cast<int>(
+          thrust::reduce(thrust::device_pointer_cast(x_data),
+                         thrust::device_pointer_cast(x_data) + x_numel,
+                         static_cast<Tx>(0), thrust::maximum<Tx>()));
+#else
+      maxlen = static_cast<int>(*std::max_element(x_data, x_data + x_numel));
+#endif
+      auto y_dim = framework::vectorize2int(x->dims());
+      y_dim.push_back(maxlen);
+      y->Resize(framework::make_ddim(y_dim));
+    }
+
+    auto out_dtype = static_cast<framework::proto::VarType::Type>(
+        ctx.Attr<int>("out_dtype"));
+    auto &dev_ctx = ctx.template device_context<DeviceContext>();
+    framework::VisitDataType(out_dtype,
+                             SequenceMaskFunctor<DeviceContext, Tx>(
+                                 dev_ctx, x_data, y, x_numel * maxlen, maxlen));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle

python/paddle/fluid/layers/nn.py

@@ -103,6 +103,7 @@ __all__ = [
     'rank_loss',
     'prelu',
     'flatten',
+    'sequence_mask',
     'stack',
 ]
 
@@ -5520,7 +5521,75 @@ def flatten(x, axis=1, name=None):
     return out
 
 
+def sequence_mask(x, maxlen=None, dtype='int64', name=None):
+    """
+    **SequenceMask Layer**
+
+    This layer outputs a mask according to the input :code:`x` and
+    :code:`maxlen` with data type of :code:`dtype`.
+
+    Supposing :code:`x` is a Tensor with shape [d_1, d_2, ..., d_n], the
+    :code:`y` is a mask with shape [d_1, d_2, ..., d_n, maxlen], where:
+    
+    .. math::
+     
+        y(i_1, i_2,..., i_n, j) = (j < x(i_1, i_2,..., i_n))
+
+    Args:
+        x (Variable): Input tensor of sequence_mask layer, 
+                      whose elements are integers less than :code:`maxlen`.
+        maxlen (int|None): Maximum length of the sequence. If :code:`maxlen`
+                           is None, it would be replace with :math:`max(x)`.
+        dtype (np.dtype|core.VarDesc.VarType|str): Data type of the output.
+        name (str|None): A name for this layer(optional). If set None, the 
+                         layer will be named automatically.  
+    
+    Returns:
+        Variable: The output sequence mask.
+    
+    """
+
+    helper = LayerHelper('sequence_mask', **locals())
+    if name is None:
+        out = helper.create_tmp_variable(dtype=dtype)
+    else:
+        out = helper.create_tmp_variable(dtype=dtype, name=name)
+
+    helper.append_op(
+        type='sequence_mask',
+        inputs={'X': [x]},
+        outputs={'Y': out},
+        attrs={
+            'max_len': maxlen if maxlen is not None else -1,
+            'out_dtype': out.dtype
+        })
+    return out
+
+
 def stack(x, axis=0):
+    """
+    **Stack Layer**
+
+    This layer stacks all of the input :code:`x` along axis.
+   
+    Input :code:`x` can be a single variable, a :code:`list` of variables, 
+    or a :code:`tuple` of variables. If :code:`x` is a :code:`list` or 
+    :code:`tuple`, the shapes of all these variables must be the same.  
+    Supposing the shape of each input is :math:`[d_0, d_1, ..., d_{n-1}]`, 
+    the shape of the output variable would be 
+    :math:`[d_0, d_1, ..., d_{axis}=len(x), ..., d_{n-1}]`. 
+    If :code:`axis` < 0, it would be replaced with :code:`axis+rank(x[0])+1`.
+    If :code:`axis` is None, it would be replaced with 0. 
+
+    Args:
+        x (Variable|list(Variable)|tuple(Variable)): Input variables. 
+        axis (int|None): The axis along which all inputs are stacked.
+    
+    Returns:
+        Variable: The stacked variable.
+    
+    """
+
     helper = LayerHelper('stack', **locals())
     axis = 0 if axis is None else axis
 

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

@@ -0,0 +1,94 @@
+# 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 op_test import OpTest
+import paddle.fluid as fluid
+from paddle.fluid.framework import convert_np_dtype_to_dtype_
+import paddle.fluid.core as core
+import numpy as np
+import copy
+import unittest
+
+
+class SequenceMaskTestBase(OpTest):
+    def initDefaultParameters(self):
+        self.op_type = 'sequence_mask'
+        self.maxlen = 10
+        self.mask_dtype = 'int64'
+        self.x = [[0, 3, 4], [5, 7, 9]]
+
+    def initParameters(self):
+        pass
+
+    def setUp(self):
+        self.initDefaultParameters()
+        self.initParameters()
+        if not isinstance(self.x, np.ndarray):
+            self.x = np.array(self.x)
+
+        self.inputs = {'X': self.x}
+        self.outputs = {'Y': self.calc_ground_truth_mask()}
+        self.attrs = {
+            'maxlen': self.maxlen,
+            'out_dtype': convert_np_dtype_to_dtype_(self.mask_dtype)
+        }
+
+    def calc_ground_truth_mask(self):
+        maxlen = np.max(self.x) if self.maxlen < 0 else self.maxlen
+        shape = self.x.shape + (maxlen, )
+        index_broadcast = np.broadcast_to(
+            np.reshape(
+                range(maxlen), newshape=[1] * self.x.ndim + [-1]),
+            shape=shape)
+        x_broadcast = np.broadcast_to(
+            np.reshape(
+                self.x, newshape=self.x.shape + (-1, )), shape=shape)
+        return (index_broadcast < x_broadcast).astype(self.mask_dtype)
+
+    def test_check_output(self):
+        self.check_output()
+
+
+class SequenceMaskTest1(SequenceMaskTestBase):
+    def initParameters(self):
+        self.mask_dtype = 'bool'
+
+
+class SequenceMaskTest2(SequenceMaskTestBase):
+    def initParameters(self):
+        self.mask_dtype = 'uint8'
+
+
+class SequenceMaskTest3(SequenceMaskTestBase):
+    def initParameters(self):
+        self.mask_dtype = 'int32'
+
+
+class SequenceMaskTest4(SequenceMaskTestBase):
+    def initParameters(self):
+        self.mask_dtype = 'float32'
+
+
+class SequenceMaskTest5(SequenceMaskTestBase):
+    def initParameters(self):
+        self.mask_dtype = 'float64'
+
+
+class SequenceMaskTest6(SequenceMaskTestBase):
+    def initParameters(self):
+        self.maxlen = -1
+
+
+if __name__ == '__main__':
+    unittest.main()