|
|
|
@ -46,47 +46,51 @@ static inline void CalcMatrixSigmaAndNormWeight(
|
|
|
|
|
Tensor* sigma, Tensor* u, Tensor* v, Tensor* weight, const int power_iters,
|
|
|
|
|
const float eps, const framework::ExecutionContext& ctx) {
|
|
|
|
|
auto& place = *ctx.template device_context<DeviceContext>().eigen_device();
|
|
|
|
|
auto blas = math::GetBlas<DeviceContext, T>(ctx);
|
|
|
|
|
auto sigma_t = EigenTensor<T, 2>::From(*sigma);
|
|
|
|
|
auto weight_t = EigenTensor<T, 2>::From(*weight);
|
|
|
|
|
auto u_t = EigenTensor<T, 1>::From(*u);
|
|
|
|
|
auto v_t = EigenTensor<T, 1>::From(*v);
|
|
|
|
|
auto u_t = EigenTensor<T, 2>::From(*u);
|
|
|
|
|
auto v_t = EigenTensor<T, 2>::From(*v);
|
|
|
|
|
|
|
|
|
|
const int h = weight->dims()[0];
|
|
|
|
|
const int w = weight->dims()[1];
|
|
|
|
|
|
|
|
|
|
Eigen::array<int, 2> perm = {1, 0};
|
|
|
|
|
Eigen::array<IndexPair, 1> product_dims = {IndexPair(1, 0)};
|
|
|
|
|
auto weight_trans_t = weight_t.shuffle(perm);
|
|
|
|
|
LOG(ERROR) << "weight: " << weight_t;
|
|
|
|
|
LOG(ERROR) << "weight_trans: " << weight_trans_t;
|
|
|
|
|
// LOG(ERROR) << "weight: " << weight_t;
|
|
|
|
|
// LOG(ERROR) << "weight_trans: " << weight_trans_t;
|
|
|
|
|
for (int i = 0; i < power_iters; i++) {
|
|
|
|
|
v_t.device(place) = weight_trans_t.contract(u_t, product_dims);
|
|
|
|
|
LOG(ERROR) << "iter v: " << v_t;
|
|
|
|
|
// v_t.device(place) = weight_trans_t.contract(u_t, product_dims);
|
|
|
|
|
blas.MatMul(*weight, true, *u, false, T(1), v, T(0));
|
|
|
|
|
// LOG(ERROR) << "iter v: " << v_t;
|
|
|
|
|
auto v_t_norm =
|
|
|
|
|
v_t.square().sum().sqrt().eval().reshape(Array1(1)).broadcast(
|
|
|
|
|
Array1(w));
|
|
|
|
|
LOG(ERROR) << "iter v_norm: " << v_t_norm;
|
|
|
|
|
// LOG(ERROR) << "iter v_norm: " << v_t_norm;
|
|
|
|
|
v_t.device(place) = v_t / (v_t_norm + v_t_norm.constant(eps));
|
|
|
|
|
LOG(ERROR) << "iter norm v: " << v_t;
|
|
|
|
|
u_t.device(place) = weight_t.contract(v_t, product_dims);
|
|
|
|
|
LOG(ERROR) << "iter u: " << u_t;
|
|
|
|
|
// LOG(ERROR) << "iter norm v: " << v_t;
|
|
|
|
|
// u_t.device(place) = weight_t.contract(v_t, product_dims);
|
|
|
|
|
blas.MatMul(*weight, false, *v, false, T(1), u, T(0));
|
|
|
|
|
// LOG(ERROR) << "iter u: " << u_t;
|
|
|
|
|
auto u_t_norm =
|
|
|
|
|
u_t.square().sum().sqrt().eval().reshape(Array1(1)).broadcast(
|
|
|
|
|
Array1(h));
|
|
|
|
|
u_t.device(place) = u_t / (u_t_norm + u_t_norm.constant(eps));
|
|
|
|
|
LOG(ERROR) << "iter norm u: " << u_t;
|
|
|
|
|
// LOG(ERROR) << "iter norm u: " << u_t;
|
|
|
|
|
}
|
|
|
|
|
LOG(ERROR) << "h" << h << "w" << w;
|
|
|
|
|
LOG(ERROR) << "u: " << u_t;
|
|
|
|
|
LOG(ERROR) << "v: " << v_t;
|
|
|
|
|
LOG(ERROR) << "weight_v: " << weight_t.contract(v_t, product_dims);
|
|
|
|
|
sigma_t.device(place) = (u_t * weight_t.contract(v_t, product_dims))
|
|
|
|
|
// LOG(ERROR) << "h" << h << "w" << w;
|
|
|
|
|
// LOG(ERROR) << "u: " << u_t;
|
|
|
|
|
// LOG(ERROR) << "v: " << v_t;
|
|
|
|
|
Tensor weight_v;
|
|
|
|
|
weight_v.mutable_data<T>({h, 1}, ctx.GetPlace());
|
|
|
|
|
blas.MatMul(*weight, false, *v, false, T(1), &weight_v, T(0));
|
|
|
|
|
auto weight_v_t = EigenTensor<T, 2>::From(weight_v);
|
|
|
|
|
// LOG(ERROR) << "weight_v: " << weight_v_t;
|
|
|
|
|
sigma_t.device(place) = (u_t * weight_v_t)
|
|
|
|
|
.sum()
|
|
|
|
|
.eval()
|
|
|
|
|
.reshape(Array2(1, 1))
|
|
|
|
|
.broadcast(Array2(h, w));
|
|
|
|
|
LOG(ERROR) << "weight: " << weight_t;
|
|
|
|
|
LOG(ERROR) << "sigma: " << sigma_t;
|
|
|
|
|
// LOG(ERROR) << "weight: " << weight_t;
|
|
|
|
|
// LOG(ERROR) << "sigma: " << sigma_t;
|
|
|
|
|
weight_t.device(place) = weight_t / sigma_t;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
@ -103,6 +107,9 @@ class SpectralNormKernel : public framework::OpKernel<T> {
|
|
|
|
|
int power_iters = ctx.Attr<int>("power_iters");
|
|
|
|
|
float eps = ctx.Attr<float>("eps");
|
|
|
|
|
|
|
|
|
|
const int h = weight->dims()[0];
|
|
|
|
|
const int w = weight->dims()[1];
|
|
|
|
|
|
|
|
|
|
Tensor weight_mat;
|
|
|
|
|
TensorCopySync(*weight, ctx.GetPlace(), &weight_mat);
|
|
|
|
|
ResizeWeight(&weight_mat, dim);
|
|
|
|
@ -113,7 +120,8 @@ class SpectralNormKernel : public framework::OpKernel<T> {
|
|
|
|
|
TensorCopySync(*u, ctx.GetPlace(), &uu);
|
|
|
|
|
TensorCopySync(*v, ctx.GetPlace(), &vv);
|
|
|
|
|
CalcMatrixSigmaAndNormWeight<DeviceContext, T>(
|
|
|
|
|
&sigma, &uu, &vv, &weight_mat, power_iters, eps, ctx);
|
|
|
|
|
&sigma, &(uu.Resize({h, 1})), &(vv.Resize({w, 1})), &weight_mat,
|
|
|
|
|
power_iters, eps, ctx);
|
|
|
|
|
TensorCopySync(weight_mat, ctx.GetPlace(), out);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|