|
|
|
@ -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 {
|
|
|
|
@ -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:
|
|
|
|
@ -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:
|
|
|
|
@ -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);
|
|
|
|
|
// ##
|
|
|
|
|
//////////////////////////////
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|