Skip to content

Commit

Permalink
refine cos-sim-op
Browse files Browse the repository at this point in the history
  • Loading branch information
@chengduoZH
chengduoZH committed Dec 16, 2017
1 parent a91efdd commit 784740d
Show file tree
Hide file tree
Showing 2 changed files with 170 additions and 74 deletions.
189 changes: 115 additions & 74 deletions paddle/operators/cos_sim_op.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#pragma once
#include "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"
#include "paddle/operators/elementwise_add_op.h"

namespace paddle {
namespace operators {
Expand All @@ -27,6 +28,28 @@ template <typename T, int MajorType = Eigen::RowMajor,
typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;

template <typename T, typename DeviceContext>
void Function_forward(T* out, T* x_norm, T* y_norm,
ElementIterator<T, DeviceContext>& x,
ElementIterator<T, DeviceContext>& y, int row, int col) {
for (int i = 0; i < row; ++i) {
T xx = 0;
T yy = 0;
T xy = 0;
for (int j = 0; j < col; ++j) {
xy += (*x) * (*y);
xx += (*x) * (*x);
yy += (*y) * (*y);
++y;
++x;
}
x_norm[i] = sqrt(xx);
y_norm[i] = sqrt(yy);

out[i] = xy / (x_norm[i] * y_norm[i]);
}
}

template <typename DeviceContext, typename T>
class CosSimKernel : public framework::OpKernel<T> {
public:
Expand All @@ -41,32 +64,63 @@ class CosSimKernel : public framework::OpKernel<T> {
out_x_norm->mutable_data<T>(context.GetPlace());
out_y_norm->mutable_data<T>(context.GetPlace());

// convert Tensor to Eigen Tensor
int rows_x = in_x->dims()[0];
int rows_y = in_y->dims()[0];
auto x = EigenMatrix<T>::Reshape(*in_x, 1);
auto y = EigenMatrix<T>::Reshape(*in_y, 1);
auto z = EigenVector<T>::Flatten(*out_z);
auto x_norm = EigenVector<T>::Flatten(*out_x_norm);
auto y_norm = EigenVector<T>::Flatten(*out_y_norm);

// compute
auto& place =
*context.template device_context<DeviceContext>().eigen_device();
auto row_along = Eigen::array<int, 1>({{1}});
x_norm.device(place) = x.square().sum(row_along).sqrt();
y_norm.device(place) = y.square().sum(row_along).sqrt();
if (rows_x == rows_y) {
auto xy = (x * y).sum(Eigen::array<int, 1>({{1}}));
z.device(place) = xy / x_norm / y_norm;
} else {
Eigen::DSizes<int, 2> bcast(rows_x, 1);
auto xy = (x * y.broadcast(bcast)).sum(row_along);
z.device(place) = xy / x_norm / y_norm.broadcast(bcast);
}

int cols = framework::product(in_x->dims()) / rows_x;
auto x_iter = ElementIterator<T, DeviceContext>(in_x->data<T>(), rows_x,
cols, rows_x, cols);
auto y_iter = ElementIterator<T, DeviceContext>(in_y->data<T>(), rows_y,
cols, rows_x, cols);

Function_forward(out_z->data<T>(), out_x_norm->data<T>(),
out_y_norm->data<T>(), x_iter, y_iter, rows_x, cols);
//
// // convert Tensor to Eigen Tensor
//// int rows_x = in_x->dims()[0];
//// int rows_y = in_y->dims()[0];
// auto x = EigenMatrix<T>::Reshape(*in_x, 1);
// auto y = EigenMatrix<T>::Reshape(*in_y, 1);
// auto z = EigenVector<T>::Flatten(*out_z);
// auto x_norm = EigenVector<T>::Flatten(*out_x_norm);
// auto y_norm = EigenVector<T>::Flatten(*out_y_norm);
//
// // compute
// auto& place =
// *context.template device_context<DeviceContext>().eigen_device();
// auto row_along = Eigen::array<int, 1>({{1}});
// x_norm.device(place) = x.square().sum(row_along).sqrt();
// y_norm.device(place) = y.square().sum(row_along).sqrt();
// if (rows_x == rows_y) {
// auto xy = (x * y).sum(Eigen::array<int, 1>({{1}}));
// z.device(place) = xy / x_norm / y_norm;
// } else {
// Eigen::DSizes<int, 2> bcast(rows_x, 1);
// auto xy = (x * y.broadcast(bcast)).sum(row_along);
// z.device(place) = xy / x_norm / y_norm.broadcast(bcast);
// }
}
};

template <typename T, typename DeviceContext>
void Function_element(T* result, ElementIterator<T, DeviceContext> dz,
ElementIterator<T, DeviceContext> y,
ElementIterator<T, DeviceContext> x_norm,
ElementIterator<T, DeviceContext> y_norm,
ElementIterator<T, DeviceContext> z,
ElementIterator<T, DeviceContext> x, int num, int block) {
for (int i = 0; i < num; ++i) {
result[i % block] += (*dz) * ((*y) / ((*x_norm) * (*y_norm)) -
(*z) * (*x) / ((*x_norm) * (*x_norm)));
++dz;
++y;
++x_norm;
++y_norm;
++z;
++x;
}
}

template <typename DeviceContext, typename T>
class CosSimGradKernel : public framework::OpKernel<T> {
public:
Expand All @@ -81,63 +135,50 @@ class CosSimGradKernel : public framework::OpKernel<T> {
auto* out_grad_y = context.Output<Tensor>(framework::GradVarName("Y"));
auto* in_grad_z = context.Input<Tensor>(framework::GradVarName("Out"));

// convert Tensor to Eigen Tensor
auto x = EigenMatrix<T>::Reshape(*in_x, 1);
auto y = EigenMatrix<T>::Reshape(*in_y, 1);
auto z = EigenMatrix<T>::Reshape(*in_z, 1);
auto x_norm = EigenMatrix<T>::Reshape(*in_x_norm, 1);
auto y_norm = EigenMatrix<T>::Reshape(*in_y_norm, 1);
auto dz = EigenMatrix<T>::Reshape(*in_grad_z, 1);

// compute gradident
int rows_x = in_x->dims()[0];
int rows_y = in_y->dims()[0];
int cols = framework::product(in_x->dims()) / rows_x;
Eigen::DSizes<int, 2> bcast_cols(1, cols);
auto z_bcast = z.broadcast(bcast_cols);
auto dz_bcast = dz.broadcast(bcast_cols);
auto x_snorm_bcast = x_norm.square().eval().broadcast(bcast_cols);
auto& place =
*context.template device_context<DeviceContext>().eigen_device();
if (rows_x == rows_y) {
auto y_snorm_bcast = y_norm.square().eval().broadcast(bcast_cols);
auto norm_prod_bcast = (x_norm * y_norm).eval().broadcast(bcast_cols);
// compute dx
if (out_grad_x) {
out_grad_x->mutable_data<T>(context.GetPlace());
auto dx = EigenMatrix<T>::Reshape(*out_grad_x, 1);
auto grad = y / norm_prod_bcast - z_bcast * x / x_snorm_bcast;
dx.device(place) = dz_bcast * grad;
}
// compute dy
if (out_grad_y) {
out_grad_y->mutable_data<T>(context.GetPlace());
auto dy = EigenMatrix<T>::Reshape(*out_grad_y, 1);
auto grad = x / norm_prod_bcast - z_bcast * y / y_snorm_bcast;
dy.device(place) = dz_bcast * grad;
}
} else {
Eigen::DSizes<int, 2> bcast_rows(rows_x, 1);
Eigen::DSizes<int, 2> bcast_rows_cols(rows_x, cols);
auto y_bcast = y.broadcast(bcast_rows);
auto y_snorm_bcast = y_norm.square().eval().broadcast(bcast_rows_cols);
auto norm_prod_bcast = (x_norm * y_norm.eval().broadcast(bcast_rows))
.eval()
.broadcast(bcast_cols);
// compute dx
if (out_grad_x) {
out_grad_x->mutable_data<T>(context.GetPlace());
auto dx = EigenMatrix<T>::Reshape(*out_grad_x, 1);
auto grad = y_bcast / norm_prod_bcast - z_bcast * x / x_snorm_bcast;
dx.device(place) = dz_bcast * grad;
}
// compute dy
if (out_grad_y) {
out_grad_y->mutable_data<T>(context.GetPlace());
auto dy = EigenVector<T>::Flatten(*out_grad_y);
auto grad = x / norm_prod_bcast - z_bcast * y_bcast / y_snorm_bcast;
dy.device(place) = (dz_bcast * grad).sum(Eigen::array<int, 1>({{0}}));
}

//////////////////////////////
// ##
auto x_iter = ElementIterator<T, DeviceContext>(in_x->data<T>(), rows_x,
cols, rows_x, cols);
auto y_iter = ElementIterator<T, DeviceContext>(in_y->data<T>(), rows_y,
cols, rows_x, cols);
auto z_iter = ElementIterator<T, DeviceContext>(in_z->data<T>(), rows_x, 1,
rows_x, cols);
auto dz_iter = ElementIterator<T, DeviceContext>(in_grad_z->data<T>(),
rows_x, 1, rows_x, cols);
auto x_norm_iter = ElementIterator<T, DeviceContext>(
in_x_norm->data<T>(), rows_x, 1, rows_x, cols);
auto y_norm_iter = ElementIterator<T, DeviceContext>(
in_y_norm->data<T>(), rows_y, 1, rows_x, cols);
// ##
//////////////////////////////
// compute dx
if (out_grad_x) {
out_grad_x->mutable_data<T>(context.GetPlace());

//////////////////////////////
// ##
Function_element(out_grad_x->data<T>(), dz_iter, y_iter, x_norm_iter,
y_norm_iter, z_iter, x_iter, rows_x * cols,
rows_x * cols);
// ##
//////////////////////////////
}
// compute dy
if (out_grad_y) {
out_grad_y->mutable_data<T>(context.GetPlace());

//////////////////////////////
// ##
Function_element(out_grad_y->data<T>(), dz_iter, x_iter, y_norm_iter,
x_norm_iter, z_iter, y_iter, rows_x * cols,
rows_y * cols);
// ##
//////////////////////////////
}
}
};
Expand Down
55 changes: 55 additions & 0 deletions paddle/operators/elementwise_op_function.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,61 @@ class MidWiseTransformIterator<T, platform::CPUDeviceContext> {
int post_;
};

template <typename T, typename Place>
class ElementIterator;

// Fixed(zcd) : Only support 2D
template <typename T>
class ElementIterator<T, platform::CPUDeviceContext> {
public:
ElementIterator(const T* ptr, int t_m, int t_n, int m, int n)
: ptr_(ptr),
index_(0),
i_(0),
j_(0),
t_m_(t_m),
t_n_(t_n),
m_(m),
n_(n) {}

ElementIterator<T, platform::CPUDeviceContext>& operator++() {
++j_;

if ((j_ == n_)) {
j_ = 0;
++i_;
}
int t_i = (t_m_ == 1) ? 0 : i_;
int t_j = (t_n_ == 1) ? 0 : j_;
index_ = t_i * t_n_ + t_j;

return *this;
}

bool operator==(
const ElementIterator<T, platform::CPUDeviceContext>& rhs) const {
return (ptr_ + index_) == &(*rhs);
}

bool operator!=(
const ElementIterator<T, platform::CPUDeviceContext>& rhs) const {
return (ptr_ + index_) != &(*rhs);
}

const T& operator*() { return ptr_[index_]; }

private:
// t_m_ == m_ || t_n_ == n_ || (t_m_ == 1 && t_m_ == 1)
const T* ptr_;
int index_;
int i_;
int j_;
int64_t t_m_;
int64_t t_n_;
int64_t m_;
int64_t n_;
};

#ifdef __NVCC__
template <typename T>
class RowwiseTransformIterator<T, platform::CUDADeviceContext>
Expand Down

0 comments on commit 784740d

Please sign in to comment.