查询TensorFlow模型推理时底层调用的BLAS例程及手动复现方法

TensorFlow推理时的BLAS调用依赖硬件和编译配置（CPU常用MKL/Eigen，GPU用cuBLAS），以下是主流场景的排查方法：

CPU场景（MKL/Eigen） #

• 开启MKL日志：运行代码前设置环境变量export MKL_VERBOSE=1（Linux/macOS）或set MKL_VERBOSE=1（Windows），MKL会在控制台输出所有调用的BLAS例程，比如单精度矩阵乘法mkl_sgemm、双精度矩阵乘法mkl_dgemm，卷积层可能调用mkl_conv系列。
• TensorFlow细粒度日志：设置export TF_CPP_VMODULE=gemm_ops=2，可打印TensorFlow内部矩阵乘法操作的细节，关联到底层BLAS调用。

GPU场景（cuBLAS） #

• 开启cuBLAS日志：设置环境变量export CUDA_VERBOSE_LOG=1，配合TF_CPP_MIN_LOG_LEVEL=0，能看到cuBLAS的调用记录，比如单精度矩阵乘cublasSgemm、批量矩阵乘cublasGemmStridedBatched（常用于批量全连接或卷积）。

高频核心BLAS例程：

• 全连接层：GEMM（通用矩阵乘法，调用最频繁）
• 卷积层：通过Im2Col转成GEMM运算，或调用专门的卷积BLAS例程（如MKL的mkl_conv_forward、cuBLAS的cublasConvolutionForward）
• 池化层：一般不调用BLAS，由TensorFlow自身实现或硬件加速指令处理

以「输入→全连接(ReLU)→全连接」的简单模型为例，步骤如下：

1. 拆解加载后的模型结构 #

先查看模型的层与参数：

import tensorflow as tf

# 加载模型
my_model = tf.saved_model.load("./saved_model")

# 遍历打印层信息
for idx, layer in enumerate(my_model.layers):
    print(f"Layer {idx}: {layer.name}, type: {type(layer).__name__}")
    if hasattr(layer, 'kernel'):
        print(f"  Kernel shape: {layer.kernel.shape}, Bias shape: {layer.bias.shape}")

2. 手动执行每一层运算 #

严格匹配模型的层逻辑计算：

import numpy as np

# 准备与predict一致的测试输入
new_input = np.random.randn(1, 784).astype(np.float32)  # 示例输入，需匹配模型输入维度

# 提取各层权重与偏置
dense1_kernel = my_model.layers[0].kernel.numpy()
dense1_bias = my_model.layers[0].bias.numpy()
dense2_kernel = my_model.layers[1].kernel.numpy()
dense2_bias = my_model.layers[1].bias.numpy()

# 手动计算第一层：矩阵乘+偏置+ReLU激活
dense1_out = np.dot(new_input, dense1_kernel) + dense1_bias
dense1_out = np.maximum(dense1_out, 0)  # ReLU

# 手动计算第二层：矩阵乘+偏置
manual_pred = np.dot(dense1_out, dense2_kernel) + dense2_bias

# 获取模型predict结果并对比
model_pred = my_model.predict(new_input)

# 检查误差（浮点精度范围内的微小差异属于正常现象）
print(f"最大绝对误差: {np.max(np.abs(model_pred - manual_pred))}")

3. 卷积模型的手动验证（可选） #

如果是卷积模型，用TensorFlow底层API模拟卷积运算：

# 提取卷积层参数
conv_layer = my_model.layers[0]
conv_kernel = conv_layer.kernel.numpy()
conv_bias = conv_layer.bias.numpy()

# 手动执行卷积运算
conv_out = tf.nn.conv2d(new_input, conv_kernel, strides=conv_layer.strides, padding=conv_layer.padding).numpy()
conv_out += conv_bias
conv_out = np.maximum(conv_out, 0)  # ReLU激活

# 后续层按全连接的方式继续计算，最终对比模型predict结果

注意：手动计算时需严格匹配模型的步长、填充、激活函数、数据类型等参数，浮点运算的微小误差（如1e-6量级）是正常的。

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