为什么 matmul 在 tensorflow 中从 1000 变为 10000 时运行得更慢？

Question

最近，我正在一个 RTX A4000 上测试混合 dtype(fp32,fp16) 中 tf.matmul 的速度。 数据集表示为 $A$，权重表示为 $B$。 matmul 是 A $times$ B 不保存结果。 我将数据放在 GPU 内存中并测试了 1000 个纪元。

让我感到困惑的是，当测试 A $times$ B 1000 个 epochs 时，需要 3.5 秒。 然而，当测试 10000 个纪元时，它应该是大约 35 秒。但结果要多得多，花费大约 10 分钟。有什么解释或意见吗？

我在这里附上独立代码：

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()

环境在这里： 张量流 2.11.0 麻木 1.23.2

硬件环境： Ubuntu 20.04.4 长期支持版 驱动版本：515.65.01
CUDA 版本：11.7

我试过不同的epoch，好像最好是1000，如果超过10000，速度会变慢

时代：100，时间成本：1.5078604221343994

时代：200，时间成本：1.7602765560150146

时代：500，时间成本：1.8624935150146484

时代：1000，时间成本：2.4284496307373047

时代：1100，时间成本：16.165714502334595

时代：1200，时间成本：36.606016635894775

时代：1500，时间成本：99.1864378452301

时代：2000，时间成本：202.15187120437622

---

单位 512

时代	时间成本	计算速度
1000	1.6995112895965576	25顶
2000	6.27071213722229	13顶
3000	10.69857382774353	12顶
5000	20.875603914260864	10顶