sequence_slice_op.h 6.1 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14 15

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
Y
Yi Wang 已提交
16 17 18
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/math/math_function.h"
#include "paddle/fluid/operators/strided_memcpy.h"
19 20 21 22 23 24 25 26 27

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using LoDTensor = framework::LoDTensor;
using LoD = framework::LoD;

template <typename T>
28
inline LoD SequenceSliceLoD(const T& in, const int64_t* offset_data,
D
dzhwinter 已提交
29
                            const int64_t* length_data) {
30
  auto out_lod = in.lod();
31 32
  size_t lod_offset = 0;

33
  auto n = in.lod()[0].size() - 1;
34 35
  out_lod[0][0] = 0;
  for (size_t i = 0; i < n; ++i) {
36
    lod_offset += length_data[i];
D
dzhwinter 已提交
37
    out_lod[0][i + 1] = lod_offset;
38 39 40 41
  }
  return out_lod;
}

Q
QI JUN 已提交
42
template <typename DeviceContext, typename T>
43
class SequenceSliceOpKernel : public framework::OpKernel<T> {
44 45 46
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    auto* in = ctx.Input<LoDTensor>("X");
47 48
    auto* offset = ctx.Input<Tensor>("Offset");
    auto* length = ctx.Input<Tensor>("Length");
49 50
    auto* out = ctx.Output<LoDTensor>("Out");

51 52 53
    auto lod = in->lod();
    auto n = lod[0].size() - 1;

D
dzhwinter 已提交
54
    PADDLE_ENFORCE_EQ(lod.size(), 1UL, "Only support one level sequence now.");
55
    PADDLE_ENFORCE_EQ(
W
wanghaox 已提交
56
        n, static_cast<size_t>(length->dims()[0]),
57
        "The size of input-sequence and length-array should be the same");
58
    PADDLE_ENFORCE_EQ(
W
wanghaox 已提交
59
        n, static_cast<size_t>(offset->dims()[0]),
60
        "The size of input-sequence and offset-array should be the same");
61

62 63
    const int64_t* offset_data = offset->data<int64_t>();
    const int64_t* length_data = length->data<int64_t>();
64 65
    framework::Tensor offset_cpu;
    framework::Tensor length_cpu;
66 67 68

    if (platform::is_gpu_place(ctx.GetPlace())) {
      offset_cpu.mutable_data<T>(offset->dims(), platform::CPUPlace());
F
fengjiayi 已提交
69
      framework::TensorCopySync(*offset, platform::CPUPlace(), &offset_cpu);
70 71 72
      offset_data = offset_cpu.data<int64_t>();

      length_cpu.mutable_data<T>(length->dims(), platform::CPUPlace());
F
fengjiayi 已提交
73
      framework::TensorCopySync(*length, platform::CPUPlace(), &length_cpu);
74 75
      length_data = length_cpu.data<int64_t>();
    }
76 77

    for (size_t i = 0; i < n; ++i) {
K
ktlichkid 已提交
78
      PADDLE_ENFORCE_LE(0, offset_data[i],
79 80 81
                        "The offset[%d] must be nonnegative.", i);
      PADDLE_ENFORCE_LE(0, length_data[i],
                        "The length[%d] must be nonnegative.", i);
K
ktlichkid 已提交
82
      PADDLE_ENFORCE_LE(lod[0][i] + offset_data[i] + length_data[i],
83
                        lod[0][i + 1], "The target tensor's length overflow.");
W
wanghaox 已提交
84
    }
85 86

    out->mutable_data<T>(ctx.GetPlace());
87
    auto out_lod = SequenceSliceLoD(*in, offset_data, length_data);
88 89 90
    auto out_dims = in->dims();
    out_dims[0] = out_lod[0][out_lod[0].size() - 1];
    out->Resize(out_dims);
91 92 93 94 95 96 97
    out->set_lod(out_lod);

    auto in_stride = framework::stride(in->dims());
    auto out_stride = framework::stride(out->dims());

    size_t out_offset = 0;
    for (size_t i = 0; i < n; ++i) {
98
      if (length_data[i] == 0) continue;
D
dzhwinter 已提交
99 100 101 102 103 104
      Tensor in_t = in->Slice(
          static_cast<int>(lod[0][i] + offset_data[i]),
          static_cast<int>(lod[0][i] + offset_data[i] + length_data[i]));

      StridedMemcpy<T>(ctx.device_context(), in_t.data<T>(), in_stride,
                       in_t.dims(), out_stride, out->data<T>() + out_offset);
105
      out_offset += length_data[i] * in_stride[0];
106 107 108 109
    }
  }
};

Q
QI JUN 已提交
110
template <typename DeviceContext, typename T>
111
class SequenceSliceGradOpKernel : public framework::OpKernel<T> {
112 113 114
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    auto* in = ctx.Input<LoDTensor>("X");
115 116
    auto* offset = ctx.Input<Tensor>("Offset");
    auto* length = ctx.Input<Tensor>("Length");
117 118 119 120 121
    auto* out_grad =
        ctx.Input<framework::LoDTensor>(framework::GradVarName("Out"));
    auto* x_grad =
        ctx.Output<framework::LoDTensor>(framework::GradVarName("X"));

122 123
    const int64_t* offset_data = offset->data<int64_t>();
    const int64_t* length_data = length->data<int64_t>();
W
wanghaox 已提交
124 125
    framework::Tensor offset_cpu;
    framework::Tensor length_cpu;
126

127 128
    if (platform::is_gpu_place(ctx.GetPlace())) {
      offset_cpu.mutable_data<T>(offset->dims(), platform::CPUPlace());
F
fengjiayi 已提交
129
      framework::TensorCopySync(*offset, platform::CPUPlace(), &offset_cpu);
130
      offset_data = offset_cpu.data<int64_t>();
131

132
      length_cpu.mutable_data<T>(length->dims(), platform::CPUPlace());
F
fengjiayi 已提交
133
      framework::TensorCopySync(*length, platform::CPUPlace(), &length_cpu);
134
      length_data = length_cpu.data<int64_t>();
135 136
    }

137
    auto lod = in->lod();
138
    auto out_lod = out_grad->lod();
139

W
wanghaox 已提交
140 141
    if (x_grad) {
      x_grad->mutable_data<T>(ctx.GetPlace());
W
wanghaox 已提交
142
      x_grad->set_lod(in->lod());
Q
QI JUN 已提交
143 144 145
      math::SetConstant<DeviceContext, T> set_zero;
      set_zero(ctx.template device_context<DeviceContext>(), x_grad,
               static_cast<T>(0));
146

W
wanghaox 已提交
147
      for (size_t i = 0; i < out_lod[0].size() - 1; ++i) {
148
        if (length_data[i] == 0) continue;
W
wanghaox 已提交
149 150 151 152
        Tensor out_grad_t =
            out_grad->Slice(static_cast<int>(out_lod[0][i]),
                            static_cast<int>(out_lod[0][i + 1]));
        auto out_grad_stride = framework::stride(out_grad_t.dims());
153

W
wanghaox 已提交
154
        auto x_grad_stride = framework::stride(x_grad->dims());
155

W
wanghaox 已提交
156 157 158
        Tensor x_grad_t = x_grad->Slice(
            static_cast<int>(lod[0][i] + offset_data[i]),
            static_cast<int>(lod[0][i] + offset_data[i] + length_data[i]));
159

W
wanghaox 已提交
160
        StridedMemcpy<T>(ctx.device_context(), out_grad_t.data<T>(),
D
dzhwinter 已提交
161 162
                         out_grad_stride, out_grad_t.dims(), x_grad_stride,
                         x_grad_t.data<T>());
W
wanghaox 已提交
163
      }
164 165 166 167 168 169
    }
  }
};

}  // namespace operators
}  // namespace paddle