我试图在tensorflow + Eigen中为我自己的op做一个元素明智的乘法运算。这是我目前使用的简单版本:
// eg) temp_res_shape = [3, 8], temp_b_shape = [1, 8]
// allocate Tensorflow tensors
Tensor temp_res;
OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<complex64>::v(),
temp_res_shape, &temp_res));
Tensor temp_a;
OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<complex64>::v(),
temp_res_shape, &temp_a));
Tensor temp_b;
OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<complex64>::v(),
temp_b_shape, &temp_b));
// These actually come from different places but the type/shape is right
// ie) I want to do this within Eigen::TensorMap if possible
auto mult_res = Tensor(temp_res).flat_inner_dims<complex64, 2>();
auto a_in = Tensor(temp_a).flat_inner_dims<complex64, 2>();
auto b_in = Tensor(temp_b).flat_inner_dims<complex64, 2>();
// convert to an array
auto a_data = a_in.data();
auto b_data = b_in.data();
auto res = mult_res.data();
for ( int i = 0; i < 3; i++ ) {
for ( int j = 0; j < 8; j++ )
res[ i*8 + j ] = a_data[ i*3 + 8 ] * b_data[j];
}
这显然是错误的做法,但我无法得到任何其他工作。我觉得它应该是这样的形式:
mult_res.device( device ) = a_in * b_in;
但这就是矩阵的倍增。我无法弄清楚如何将b_in转换为对角矩阵以便以这种方式相乘:/ 我觉得这应该是微不足道的,但我无法解决(我的cpp不是很好)。提前谢谢!