squared_l2_distance_op.h 4.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

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 "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
23 24 25
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;
26 27
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
28
using EigenMatrix = framework::EigenMatrix<T, MajorType, IndexType>;
29 30 31 32 33

template <typename Place, typename T>
class SquaredL2DistanceKernel : public framework::OpKernel {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
    auto* in0 = context.Input<Tensor>("X");
    auto* in1 = context.Input<Tensor>("Y");
    auto* out0 = context.Output<Tensor>("sub_result");
    auto* out1 = context.Output<Tensor>("Out");

    auto in0_dims = in0->dims();
    auto in1_dims = in1->dims();

    int cols = framework::product(in0_dims) / in0_dims[0];
    // reduce dimensions except the first
    auto x =
        EigenMatrix<T>::From(*in0, framework::make_ddim({in0_dims[0], cols}));
    auto y =
        EigenMatrix<T>::From(*in1, framework::make_ddim({in1_dims[0], cols}));

    out0->mutable_data<T>(context.GetPlace());
    out1->mutable_data<T>(context.GetPlace());
    auto sub_result = EigenMatrix<T>::From(*out0);
52
    auto z = EigenVector<T>::Flatten(*out1);
53

54
    auto place = context.GetEigenDevice<Place>();
55 56
    auto x_dims = x.dimensions();
    auto y_dims = y.dimensions();
57
    // buffer the substraction result
58
    if (y_dims[0] == 1 && x_dims[0] > y_dims[0]) {
Y
yangyaming 已提交
59 60
      sub_result.device(place) =
          x -
61
          y.broadcast(Eigen::array<int, 2>({{static_cast<int>(x_dims[0]), 1}}));
62
    } else {
63
      sub_result.device(place) = x - y;
64
    }
Y
yangyaming 已提交
65
    auto sub_res_pow2 = sub_result * sub_result;
66
    z.device(place) = sub_res_pow2.sum(Eigen::array<int, 1>({{1}}));
67 68 69 70 71 72 73
  }
};

template <typename Place, typename T>
class SquaredL2DistanceGradKernel : public framework::OpKernel {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
74 75 76 77
    auto* in0 = context.Input<Tensor>("sub_result");
    auto* in1 = context.Input<Tensor>(framework::GradVarName("Out"));
    auto* x_g = context.Output<Tensor>(framework::GradVarName("X"));
    auto* y_g = context.Output<Tensor>(framework::GradVarName("Y"));
78

79 80
    auto sub_result = EigenMatrix<T>::From(*in0);
    auto out_grad = EigenMatrix<T>::From(*in1);
81

82 83
    auto x_dims = x_g->dims();
    auto y_dims = y_g->dims();
84

85
    int cols = framework::product(x_dims) / x_dims[0];
86
    // calculate gradient
87 88 89
    auto grad_mat = 2 *
                    (out_grad.broadcast(Eigen::array<int, 2>({{1, cols}}))) *
                    sub_result;
90 91

    // propagate back to input
92
    auto eigen_place = context.GetEigenDevice<Place>();
Y
yangyaming 已提交
93
    if (x_g) {
94 95 96 97
      x_g->mutable_data<T>(context.GetPlace());
      // eigen matrix
      auto x_grad =
          EigenMatrix<T>::From(*x_g, framework::make_ddim({x_dims[0], cols}));
98
      // dimensions are same with subResult
99
      x_grad.device(eigen_place) = grad_mat;
100
    }
101

Y
yangyaming 已提交
102
    if (y_g) {
103 104
      y_g->mutable_data<T>(context.GetPlace());

Y
yangyaming 已提交
105 106 107
      PADDLE_ENFORCE_GE(sub_result.dimensions()[0], y_dims[0],
                        "First dimension of gradient must be greater or "
                        "equal than first dimension of target.");
108

109
      if (sub_result.dimensions()[0] == y_dims[0]) {
110 111
        auto y_grad =
            EigenMatrix<T>::From(*y_g, framework::make_ddim({y_dims[0], cols}));
112
        y_grad.device(eigen_place) = -1 * grad_mat;
113
      } else {
114 115
        auto col_sum_res = -1 * (grad_mat.sum(Eigen::array<int, 1>({{0}})));
        auto y_grad = EigenVector<T>::Flatten(*y_g);
116
        y_grad.device(eigen_place) = col_sum_res;
117 118
      }
    }
119 120 121 122 123
  }
};

}  // namespace operators
}  // namespace paddle