0

矩阵链是矩阵矩阵乘积的链。我考虑以下矩阵链：

美国广播公司；其中A和B的大小为 3000x3000，C的大小为 3000x600

有两种方法可以评估上述表达式，它们在性能上存在显着差异：

变体 1：(AB)C：6.48e10 FLOPs

变体 2: A(BC) : 2.16e10 FLOPs

矩阵矩阵乘法AB的成本是2xmxnxk，其中A的大小为mxn，B的大小为nxk。使用这个公式，我获得了上述变体的 FLOP 性能。

如果括号没有明确指定，我的 TensorFlow 构建（版本 2.8，禁用 Eager 模式）仅选择变体 1（从左到右括号），其执行时间几乎是变体 2 的三倍。虽然我可以优化这个并且通过显式计算矩阵乘法的 FLOP 手动加括号，我很好奇这是否可以由 TensorFlow 使用的 Grappler 图优化器自动完成？是否有任何其他图形优化器可以自动选择最佳括号？

观察不同括号的性能效果的示例脚本


import tensorflow as tf
import os
import time

class bcolors:
    WARNING = '\033[93m'
    ENDC = '\033[0m'


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


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

#Problem size
n = 3000
reps = 10
DTYPE = tf.float32


@tf.function
def mc_non_optimized(A,B,C):
    # Default Parenthesization (Variant 1)
    start =  tf.timestamp()
    with tf.control_dependencies([start]):
        ret = A@B@C
    with tf.control_dependencies([ret]):
        end =  tf.timestamp()
        tf.print("Non Optimized : ", end-start)
    
    return ret

@tf.function
def mc_optimized(A,B,C):
    #Optimized parenthesization (Variant 2)
    start =  tf.timestamp()
    with tf.control_dependencies([start]):
        # I do not want to manually find the optimum parethesization every time
        ret = A@(B@C)
    with tf.control_dependencies([ret]):
        end =  tf.timestamp()
        tf.print("Optimized : ", end-start)

    
    return ret


A = tf.random.normal([n, n], dtype=DTYPE)
B = tf.random.normal([n, n], dtype=DTYPE)
C = tf.random.normal([n, int(n/5)], dtype=DTYPE)


for i in range(reps):
   ret = mc_non_optimized(A,B,C)
   ret = mc_optimized(A,B,C)
   tf.print("\n")

使用 Intel MKL 和 Python 3.9.7 构建的 TensorFlow 2.8 (CPU) 的执行时间在 Mac book pro 2018 Big sur 上运行

  • 变体 1（默认括号）：0.65s
  • 变体 2（优化括号）：0.2s
4

0 回答 0