TensorFlow中tf.matmul随Epoch增加速度骤降的原因咨询

在RTX A4000显卡上测试混合精度（fp32、fp16）下tf.matmul的运算速度，将数据集$A$与权重$B$置于GPU显存中执行$A×B$运算且不保存结果。测试1000个Epoch耗时约3.5s，但测试10000个Epoch时，预期耗时约35s，实际却花费约10分钟。

测试不同Epoch数的耗时如下：

• epoch:100，耗时:1.5078604221343994s
• epoch:200，耗时:1.7602765560150146s
• epoch:500，耗时:1.8624935150146484s
• epoch:1000，耗时:2.4284496307373047s
• epoch:1100，耗时:16.165714502334595s
• epoch:1200，耗时:36.606016635894775s
• epoch:1500，耗时:99.1864378452301s
• epoch:2000，耗时:202.15187120437622s

补充units=512时的测试数据：

完整测试代码：

import time
import numpy as np
import tensorflow as tf

@tf.function
def inference(x, y):
    return tf.matmul(x,y) 

def _test_mnist_fp32():
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = np.float32(x_train/1024)
    data_size = 60000*28*28
    epoch = 1000
  
    # data_holder = tf.Variable(initial_value=tf.reshape(x_train, shape=(60000,-1)))
    data_holder = tf.constant(tf.reshape(x_train, shape=(60000,-1)))
    units = 51200
    # weights = tf.Variable(initial_value=tf.random.normal(shape=(28*28, units)))
    weights = tf.constant(tf.random.normal(shape=(28*28, units)))

    start_time = time.time()
    for i in range(epoch):
        inference(data_holder, weights)
    print(data_holder.device)
    end_time = time.time()
    
    time_cost = end_time-start_time
    Mulplus = 2*data_size*units*epoch/1024/1024/1024/1024
    print("time cost:{}".format(time_cost))
    print("total Mulplus: {}Tops, compute speed:{}Tops/s".format(Mulplus, int(Mulplus/time_cost)))

if __name__ == "__main__":
    _test_mnist_fp32()

环境信息：

• 软件环境：tensorflow 2.11.0，numpy 1.23.2
• 硬件环境：Ubuntu 20.04.4 LTS，Driver Version: 515.65.01，CUDA Version: 11.7

1. 计算图节点累积：@tf.function装饰的函数会构建TensorFlow计算图，Python循环中反复调用inference会不断向计算图中添加重复的matmul节点。当Epoch数超过阈值后，计算图规模急剧膨胀，GPU调度和执行的开销大幅增加，导致运算速度下降。
2. GPU显存碎片化：即使不保存运算结果，大量重复的矩阵乘法产生的中间张量可能无法被及时回收，随着Epoch数增加，显存碎片增多，GPU需要花费更多时间进行内存分配与回收，拖慢运算效率。
3. TensorFlow调度机制过载：小批量Epoch时，TensorFlow可合并优化运算任务；当Epoch数过大，调度器负载增加，无法有效优化任务调度，单步运算延迟上升。
4. CUDA资源耗尽：大量重复运算可能耗尽CUDA上下文的流、事件等资源，后续运算需要等待资源释放，产生额外等待时间。

1. 将循环逻辑放入tf.function内部：避免Python循环中反复调用图函数，减少节点重复构建，示例修改如下：

@tf.function
def run_epochs(data, weights, epoch):
    for i in tf.range(epoch):
        tf.matmul(data, weights)

# 测试函数中调用
start_time = time.time()
run_epochs(data_holder, weights, epoch)
end_time = time.time()

2. 启用显存按需分配：开启TensorFlow显存动态增长，减少显存碎片：

tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)

3. 使用TensorFlow原生循环API：用tf.while_loop替代Python循环，让TensorFlow更好地优化循环内运算，降低调度开销。
4. 合并运算减少循环次数：若仅测试运算速度，可构造更大张量完成等价批量运算，避免循环带来的额外开销。
5. 升级TensorFlow版本：TensorFlow 2.11为旧版本，后续版本对GPU内存管理和图优化有较多改进，升级到2.15+或最新稳定版可能缓解问题。
6. 监控显存状态：用nvidia-smi实时查看GPU显存占用，确认Epoch数增加时是否出现显存占用急剧上升，验证内存累积假设。

内容的提问来源于stack exchange，提问作者Enzhi