张量流中 SpatialConvolution 实现的性能

Question

参考tensorflow的实现（使用eigen）实现了一个SpatialConvolution函数。 tensorflow中的实现位于SpatialConvolution and I also find one related reply about the implementation :

我的实现如下：（由于我的数据是行优先的，所以我只保留了一半的代码）

// Description: Convolution                                                                          
// Input:                                                                                           
//      - name: input0     type: float     shape: Shape{7680, 15, 200, 1}                           
//      - name: input1     type: float     shape: Shape{5, 200, 1, 200}                             
// Output:                                                                                          
//      - name: output0    type: float     shape: Shape{7680, 11, 1, 200}
void Convolution_float_float_float_cpu_Convolution_270(float* input0, float* input1, float* output0)
{

Eigen::array<Eigen::IndexPair<Eigen::Index>, 1> contract_dims;
contract_dims[0] = Eigen::IndexPair<Eigen::Index>(1, 0);

Eigen::array<Eigen::Index, 4> in_dims({7680, 15, 200, 1});
Eigen::array<Eigen::Index, 4> out_dims({7680, 11, 1, 200});
Eigen::array<Eigen::Index, 4> kernel_dims({5, 200, 1, 200});
Eigen::DSizes<Eigen::Index, 2> pre_contract_dims;
pre_contract_dims[1] = kernel_dims[2] * kernel_dims[1] * kernel_dims[0];
pre_contract_dims[0] = out_dims[1] * out_dims[2];
for (int i = 0; i < 1; ++i) {
  pre_contract_dims[0] *= in_dims[i];
}

Eigen::DSizes<Eigen::Index, 4> post_contract_dims;
post_contract_dims[3] = kernel_dims[3];
post_contract_dims[2] = out_dims[2];
post_contract_dims[1] = out_dims[1];
for (int i = 0; i < 1; ++i) {
  post_contract_dims[i] = in_dims[i];
}

Eigen::DSizes<Eigen::Index, 2> new_kernel_dims;
new_kernel_dims[0] = kernel_dims[2] * kernel_dims[1] * kernel_dims[0];
new_kernel_dims[1] = kernel_dims[3];

Eigen::TensorMap<Eigen::Tensor<float, 4, Eigen::RowMajor>>
    in(static_cast<float *>(input0), in_dims),
    out(static_cast<float *>(output0), out_dims),
    kernel(static_cast<float *>(input1), kernel_dims);

out.device(*global_thread_pool_device) = in
    .extract_image_patches(kernel_dims[1], kernel_dims[0], 1,
                           1, 1, 1,
                           Eigen::PADDING_VALID)
    .reshape(pre_contract_dims)
    .contract(kernel.reshape(new_kernel_dims), contract_dims)
    .reshape(post_contract_dims);
}

通过处理相同的数据并将线程池中的线程数设置为 1（intra_op_parallelism_threads 在 tensorflow 中），看起来我的实现比 tensorflow 慢了大约 30%。我的编译器选项是“-std=gnu++11 -O3 -march=native”，没有启用tensorflow的XLA。我不知道是什么导致了性能差距。如果有人可以提供一些提示，那将是一个很大的帮助。

Answer 1

深入研究代码后，我们发现tensorflow基于MKL实现了自定义特征核。通过 eigen_contraction_kernel.h/.cpp 和 eigen_spatial_convolutions.h 中的实现，我们可以获得与 tensorflow 相同的性能。

张量流中 SpatialConvolution 实现的性能

Performance of SpatialConvolution implementation in tensorflow

c++

tensorflow

convolution