为什么我的 CPU Float64 tf.matmul 在 TensorFlow2 中的性能明显慢于 NumPy matmul，即使在图形模式下也是如此？

Question

我正在比较 TensorFlow 2 和 NumPy 中矩阵-矩阵产品的单线程性能。我分别比较了单精度 (float32) 和双精度 (float64)。我发现 NumPy 在单精度和双精度（DGEMM 和 SGEMM）方面的性能几乎等同于英特尔 MKL C++ 实现（用作矩阵乘法的基准）。但在TensorFlow中，只有单精度（float32）性能与MKL相当，双精度（float64）性能明显较慢。 为什么使用双精度数据时 Tensorflow 变慢？

示例脚本：

我认为以下实例可以重现我的观察结果。考虑矩阵乘法：

C = AB where A and B are of size 3000x3000

TensorFlow2和NumPy代码如下：

Tensorflow2代码

import tensorflow as tf
import os
import time


#Check if MKL is enabled
import tensorflow.python.framework as tff
print("MKL Enabled : ", tff.test_util.IsMklEnabled())


#Set threads
tf.config.threading.set_inter_op_parallelism_threads(1)
tf.config.threading.set_intra_op_parallelism_threads(1)

#Problem size
N = 3000
REPS = 20
DTYPE = tf.float64
#DTYPE = tf.float32


@tf.function
def gemm_implicit_noup(A, B):
    #C = A @ B
    start = tf.timestamp()
    with tf.control_dependencies([start]):
        C = tf.matmul(A,B)
    with tf.control_dependencies([C]):
        end = tf.timestamp()
    tf.print(end-start)
    return C

tf.config.run_functions_eagerly(False)

A = tf.random.normal([N, N], dtype=DTYPE)
B = tf.random.normal([N, N], dtype=DTYPE)


#Building Trace
C = gemm_implicit_noup(A,B)

for i in range(REPS):
   C = gemm_implicit_noup(A,B)

Numpy 代码

import os
os.environ["OMP_NUM_THREADS"] = "1"
import numpy as np
import time

N = 3000
REPS = 20
DTYPE = np.float64
#DTYPE = np.float32

def gemm_implicit_noup(A, B):
    #C = A @ B
    C = np.matmul(A,B)
    return C



A = np.random.randn(N,N).astype(DTYPE)
B = np.random.randn(N,N).astype(DTYPE)

for i in range(REPS):
   start = time.perf_counter()
   C = gemm_implicit_noup(A,B)
   end = time.perf_counter()
   print(end-start)

系统和安装设置：

性能是在 Intel Xeon Skylake 2.1 GHz 和 CentOS 7 以及 MacBook Pro 2018 和 BigSur 上进行比较的。在使用英特尔 MKL 构建的 Tensorflow 2.7 和 2.8 上比较了性能。检查了 Python 3.9.7 和 3.7.4。我比较了单线程性能，以便可以可靠地重现结果。我在所有设置中观察到类似的性能数字：

单精度性能符合预期：

英特尔 MKL C++ SGEMM ~ 0.5s
NumPy float32 ~ 0.5s
TensorFlow float32 ~ 0.5s

但是双精度性能：

英特尔 MKL C++ DGEMM ~ 0.9s
NumPy float64 ~ 1s
TensorFlow float64 > 2.5s（慢得多！！）

Answer 1

假设您使用的是 Intel® AVX-512 指令支持处理器，请尝试安装 Intel® Optimization for TensorFlow Wheel通过 PIP 专门为 AVX512 构建。这些软件包在英特尔® 网站上以 *.whl 形式提供，适用于特定 Python 版本，或者可以使用以下命令安装 Python 版本 3.7、3.8 和 3.9（仅限 Linux） .

pip install intel-tensorflow-avx512==2.7.0

这在英特尔® 官方网站及其子部分的以下链接中有记录：

Intel® Optimization for TensorFlow: Installation Guide

Intel® Optimization for TensorFlow: Install the Intel® Optimization for TensorFlow Wheel via PIP

AVX512 是一个 Single Instruction Multiple Data (SIMD) instruction set specifically designed to handle complex data types like double-precision numbers. In order to take full advantage of Intel® architecture and to extract maximum performance, the TensorFlow framework has been optimized using oneAPI Deep Neural Network Library (oneDNN) 原语，一个流行的深度学习应用性能库。作为额外的优化步骤，还可以尝试在运行 TensorFlow 代码之前使用以下命令在 Linux 终端内将环境变量 TF_ENABLE_ONEDNN_OPTS 设置为 1 :

export TF_ENABLE_ONEDNN_OPTS=1

使用您提供的代码得到的双精度矩阵乘积的单线程性能如下。此测试是在 Intel® Xeon® Platinum 8260M CPU @ 2.40GHz 和 Python 3.8 以及 英特尔® MKL 和 AVX512 优化的 TensorFlow 2.7.

NumPy float64 ~ 1.44s
TensorFlow float64（启用 MKL）~ 2.77s
TensorFlow float64（启用 MKL，优化 AVX512，启用 oneDNN 优化）~ 1.19s

为什么我的 CPU Float64 tf.matmul 在 TensorFlow2 中的性能明显慢于 NumPy matmul，即使在图形模式下也是如此？

Why is my CPU Performance of Float64 tf.matmul in TensorFlow2 significantly slower than the NumPy matmul, even in the graph mode?

performance

numpy

matrix-multiplication

intel-mkl

tensorflow