You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Paddle/paddle/fluid/operators/sequence_ops/sequence_concat_op.h

133 lines
4.5 KiB

// 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
#include <vector>
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/detail/safe_ref.h"
#include "paddle/fluid/operators/math/concat_and_split.h"
namespace paddle {
namespace operators {
namespace detail {
template <typename Container>
inline framework::LoD ConcatLoD(const Container &xs,
std::vector<framework::Tensor> *xs_in_order) {
std::vector<size_t> result;
result.resize(xs[0].get().lod()[0].size());
for (size_t i = 1; i < result.size(); ++i) {
size_t sum = 0;
for (size_t j = 0; j < xs.size(); ++j) {
auto &x_lod = xs[j].get().lod()[0];
const framework::Tensor &tensor = xs[j].get();
xs_in_order->emplace_back(tensor.Slice(x_lod[i - 1], x_lod[i]));
sum += x_lod[i];
}
result[i] = sum;
}
framework::LoD lod;
lod.emplace_back(result);
return lod;
}
} // namespace detail
template <typename DeviceContext, typename T>
class SeqConcatKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext &context) const override {
auto xs = detail::VectorRef(context.MultiInput<framework::LoDTensor>("X"),
"Cannot find multiple input X");
auto &out = detail::Ref(context.Output<framework::LoDTensor>("Out"),
"Cannot find output");
size_t lod_size = 0;
for (auto &x : xs) {
if (lod_size == 0) {
lod_size = x.get().lod()[0].size();
} else {
PADDLE_ENFORCE_EQ(
lod_size, x.get().lod()[0].size(),
"The number of sequence must be same between each input");
}
}
PADDLE_ENFORCE_NE(lod_size, 0, "Each input must have sequence information");
std::vector<framework::Tensor> x_in_order;
out.set_lod(detail::ConcatLoD(xs, &x_in_order));
out.mutable_data<T>(context.GetPlace());
math::ConcatFunctor<DeviceContext, T> functor;
functor(context.template device_context<DeviceContext>(), x_in_order, 0,
&out);
}
};
template <typename DeviceContext, typename T>
class SeqConcatGradKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext &context) const override {
auto xs = context.MultiInput<framework::LoDTensor>("X");
auto dxs =
context.MultiOutput<framework::LoDTensor>(framework::GradVarName("X"));
PADDLE_ENFORCE_EQ(xs.size(), dxs.size());
for (size_t i = 0; i < dxs.size(); ++i) {
if (dxs[i] != nullptr) {
dxs[i]->set_lod(xs[i]->lod());
dxs[i]->mutable_data<T>(context.GetPlace());
}
}
std::vector<framework::Tensor> sliced_x;
std::vector<boost::variant<boost::blank, framework::Tensor>> sliced_dx;
for (size_t i = 1; i < xs[0]->lod()[0].size(); ++i) {
for (size_t j = 0; j < xs.size(); ++j) {
const framework::LoDTensor *x = xs[j];
framework::LoDTensor *dx = dxs[j];
auto &x_lod = x->lod()[0];
sliced_x.emplace_back(x->Slice(x_lod[i - 1], x_lod[i]));
if (dx != nullptr) {
sliced_dx.emplace_back(dx->Slice(x_lod[i - 1], x_lod[i]));
} else {
sliced_dx.emplace_back(boost::blank());
}
}
}
math::SplitFunctor<DeviceContext, T> functor;
std::vector<const framework::Tensor *> sliced_x_ptr;
std::vector<framework::Tensor *> sliced_dx_ptr;
for (auto &x : sliced_x) {
sliced_x_ptr.emplace_back(&x);
}
for (auto &dx : sliced_dx) {
try {
sliced_dx_ptr.emplace_back(&boost::get<framework::Tensor>(dx));
} catch (boost::bad_get &) {
sliced_dx_ptr.emplace_back(nullptr);
}
}
functor(context.template device_context<DeviceContext>(),
detail::Ref(
context.Input<framework::Tensor>(framework::GradVarName("Out")),
"Sequence Concat OG must be set"),
sliced_x_ptr, 0, &sliced_dx_ptr);
}
};
} // namespace operators
} // namespace paddle