为何Numpy实现的神经网络比Keras快近25倍？

测试背景 #

我正在对比Numpy自定义神经网络和Keras实现的神经网络的性能，使用TensorFlow 2.13.0版本，已通过tf.compat.v1.disable_eager_execution()禁用即时执行，在Jupyter Notebook中开展基准测试。

Keras模型实现 #

class TFModel():
    def __init__(self,input_shape,outputs,hidden_units,**kwargs):
        _input = Input(shape=(input_shape))
        layer = Dense(units = input_shape,activation='relu',use_bias=False)(_input)
        for hu in hidden_units:
            layer = Dense(units = hu,activation='relu',use_bias=False)(layer)
        outputs = [Dense(units=out,activation='sigmoid',use_bias=False)(layer) for out in outputs]
        self.model = Model(inputs = _input,outputs=outputs)
        self.total_weights = self.model.count_params()
        
    def predict(self,x):
        prediction = self.model.predict(x)
        return prediction

Numpy自定义神经网络实现 #

class NeuralNetwork(ActivationFunctions):
    def __init__(self,layers,seed=None):
        self.layers = {}
        self.total_weights = 0
        self._module_to_use = np.random
        self.seed = seed
        if seed is not None:
            self._module_to_use = self._generate_random_state()
        inputs = layers[:-1]
        outputs = layers[1:]
        
        for ins,outs in zip(inputs,outputs):
            self.total_weights += sum(ins.nodes)*sum(outs.nodes)
            layers = self._generate_wt(ins.nodes,outs.nodes)
            actual_len = len(self.layers)
            for idx,layer in enumerate(layers):
                activation = outs.activation                
                if actual_len not in self.layers:
                    self.layers[actual_len] = {}
                self.layers[actual_len][idx] = {
                    'layer':layer,
                    'activation':activation
                }
        self.layers[actual_len]['type'] = 'output'
    
    def _generate_random_state(self):
        return RandomState(MT19937(SeedSequence(self.seed)))
    
    def _generate_wt(self,x,y):
        return [self._module_to_use.random(size=(_x,_y)) for _x,_y in product(x,y)]

    def f_forward(self,x):
        # hidden
        output = []
        for layer in self.layers.values():
            is_output = layer.get('type',False)
            for lay in layer.values():
                if lay == 'output':
                    continue
                activation = getattr(self,lay['activation'])
                lay = lay['layer']
                z = x.dot(lay)
                z = activation(z)
                if is_output:
                    output.append(z)
            x = z
        return output

辅助类实现 #

class ActivationFunctions():
    def softmax(self,x):
        """Compute softmax values for each sets of scores in x."""
        e_x = np.exp(x - np.max(x))
        return e_x / e_x.sum()
    
    def relu(self,x):
        return x * (x>0)
    
    def sigmoid(self,x):
        return(1/(1 + np.exp(-x)))

class Layer():
    def __init__(self,nodes,activation):
        if isinstance(nodes,int):
            nodes = (nodes,)
        self.nodes = nodes
        self.activation = activation

测试输入 #

i = np.random.randint(0,2,size=(1,24))

测试结果 #

Keras模型测试 #

input_shape = 24
outputs = [6,2]
hidden_units = [10]

tf_model = TFModel(input_shape,outputs,hidden_units)

%timeit -n 100 -r 10 tf_model.model.predict(i)
>416 µs ± 19.1 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

Numpy模型测试 #

layers = [(24,24,10,(6,2)),(None,'relu','relu','softmax')]
layers = [Layer(node,act) for node,act in zip(*layers)]
nn =  NeuralNetwork(layers)

%timeit -n 100 -r 10 nn.f_forward(i)
>17 µs ± 4.27 µs per loop (mean ± std. dev. of 10 runs, 100 loops each)

差距原因与操作校验 #

核心性能差距原因 #

1. model.predict()的额外开销：Keras的predict()方法并非仅执行前向计算，还包含输入格式校验、批次处理逻辑、TensorFlow内部张量转换、输出格式整理等操作。在单样本输入场景下，这些额外开销占比远高于实际计算成本。

2. TensorFlow图模式的固定开销：即使禁用即时执行，静态图模式仍存在图初始化、节点调度的固定成本。而Numpy代码直接调用底层BLAS/LAPACK做矩阵运算，没有框架层面的调度损耗。

3. 测试场景的局限性：单样本推理无法体现TensorFlow的并行计算优势。当输入样本量增大（比如(1000,24)），框架固定开销占比会骤降，TensorFlow的多线程/硬件加速能力会反超Numpy的单线程运算。

操作校验与优化建议 #

• 模型结构对齐检查：Keras输出层用sigmoid激活，Numpy模型用softmax，虽然不影响推理速度，但需确认层数、节点数完全匹配；Keras默认用glorot_uniform初始化权重，Numpy用np.random.random()，初始化方式不同但不影响推理性能。
• 优化Keras测试方式：跳过predict()的额外逻辑，直接调用模型或用tf.function编译前向计算，得到更真实的计算耗时：

  # 直接调用模型
  %timeit -n 100 -r 10 tf_model.model(i)
  
  # 用tf.function编译
  @tf.function
  def tf_predict(x):
      return tf_model.model(x)
  %timeit -n 100 -r 10 tf_predict(i)
  

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