Производительность реализации SpatialConvolution в тензорном потоке - PullRequest
Новая
       66

Производительность реализации SpatialConvolution в тензорном потоке

0 голосов
/ 07 мая 2020

Я реализовал функцию SpatialConvolution, сославшись на реализацию tenorflow (используйте собственный). Реализация в тензорном потоке находится по адресу SpatialConvolution , и я также нахожу один связанный ответ о реализации: { ссылка }

Моя реализация выглядит следующим образом: (поскольку мои данные является строкой, я сохраняю только половину кода) 

// 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 в тензорном потоке), похоже, моя реализация составляет около 30% медленнее, чем тензорный поток. Моя опция компилятора - "-std = gnu ++ 11 -O3 -march = native", а XLA тензорного потока не включен. Я понятия не имею, что привело к разрыву в производительности. Если бы кто-нибудь мог дать несколько подсказок, это было бы большим подспорьем.

1 Ответ

0 голосов
/ 26 мая 2020

Покопавшись в коде, мы обнаружили, что в тензорном потоке реализованы собственные ядра на основе MKL. С реализациями в eigen_contraction_kernel.h /.cpp и eigen_spatial_convolutions.h, мы можем получить ту же производительность с тензорным потоком.

...