|
|
|
@ -21,6 +21,9 @@ namespace operators {
|
|
|
|
|
|
|
|
|
|
using Tensor = framework::Tensor;
|
|
|
|
|
using LoDTensor = framework::LoDTensor;
|
|
|
|
|
template <typename T, int MajorType = Eigen::RowMajor,
|
|
|
|
|
typename IndexType = Eigen::DenseIndex>
|
|
|
|
|
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;
|
|
|
|
|
template <typename T, int MajorType = Eigen::RowMajor,
|
|
|
|
|
typename IndexType = Eigen::DenseIndex>
|
|
|
|
|
using EigenMatrix = framework::EigenMatrix<T, MajorType, IndexType>;
|
|
|
|
@ -43,8 +46,8 @@ class SequenceAvgPoolKernel : public framework::OpKernel {
|
|
|
|
|
static_cast<int>(lod[0][i + 1]));
|
|
|
|
|
Tensor out_t = out->Slice<T>(i, i + 1);
|
|
|
|
|
int64_t h = static_cast<int64_t>(lod[0][i + 1] - lod[0][i]);
|
|
|
|
|
auto in_e = EigenMatrix<T>::From(in_t, {h, w});
|
|
|
|
|
auto out_e = EigenMatrix<T>::From(out_t, {h, w});
|
|
|
|
|
auto in_e = EigenMatrix<T>::From(in_t, framework::make_ddim({h, w}));
|
|
|
|
|
auto out_e = EigenVector<T>::Flatten(out_t);
|
|
|
|
|
out_e.device(place) = in_e.mean(Eigen::array<int, 1>({{0}}));
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
@ -54,9 +57,9 @@ template <typename Place, typename T>
|
|
|
|
|
class SequenceAvgPoolGradKernel : public framework::OpKernel {
|
|
|
|
|
public:
|
|
|
|
|
void Compute(const framework::ExecutionContext& context) const override {
|
|
|
|
|
auto* in = context.Output<LoDTensor>("X");
|
|
|
|
|
auto* in_g = context.Output<LoDTensor>(framework::GradVarName("X"));
|
|
|
|
|
auto* in = context.Input<LoDTensor>("X");
|
|
|
|
|
auto* out_g = context.Input<LoDTensor>(framework::GradVarName("Out"));
|
|
|
|
|
auto* in_g = context.Output<LoDTensor>(framework::GradVarName("X"));
|
|
|
|
|
|
|
|
|
|
auto dims = in->dims();
|
|
|
|
|
auto lod = in->lod();
|
|
|
|
@ -71,7 +74,7 @@ class SequenceAvgPoolGradKernel : public framework::OpKernel {
|
|
|
|
|
int64_t h = static_cast<int64_t>(lod[0][i + 1] - lod[0][i]);
|
|
|
|
|
auto in_g_e = EigenMatrix<T>::From(in_g_t, {h, w});
|
|
|
|
|
auto out_g_e = EigenMatrix<T>::From(out_g_t, {1, w});
|
|
|
|
|
Eigen::DSizes<int, 2> bcast(h, w);
|
|
|
|
|
Eigen::DSizes<int, 2> bcast(h, 1);
|
|
|
|
|
in_g_e.device(place) = (out_g_e / static_cast<T>(h)).broadcast(bcast);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|