TensorFlow2中替换Conv2D/矩阵乘法实现并保留默认梯度

核心思路是通过tf.custom_gradient装饰器，将前向传播用你的C实现（CTypes调用），反向传播复用TensorFlow原生运算的梯度逻辑，这样既替换了前向计算，又保留了默认梯度，同时满足CTypes调用的要求。

一、矩阵乘法（替换tf.matmul） #

1. 封装带梯度的自定义matmul #

import tensorflow as tf
import numpy as np
import ctypes

# 假设已通过CTypes加载C库，示例：c_approx = ctypes.CDLL("./your_custom_ops.so")
# 定义C函数的参数与返回值类型（需匹配你的C代码接口）
c_approx.custom_matmul.argtypes = [np.ctypeslib.ndpointer(dtype=np.float32),
                                  np.ctypeslib.ndpointer(dtype=np.float32),
                                  ctypes.c_int, ctypes.c_int, ctypes.c_int,
                                  np.ctypeslib.ndpointer(dtype=np.float32)]
c_approx.custom_matmul.restype = None

@tf.custom_gradient
def custom_matmul(a, b, transpose_a=False, transpose_b=False):
    # 前向传播：调用C代码
    a_np = tf.convert_to_tensor(a).numpy()
    b_np = tf.convert_to_tensor(b).numpy()
    
    # 处理转置参数（可根据C代码能力选择在Python或C侧处理）
    if transpose_a:
        a_np = a_np.T
    if transpose_b:
        b_np = b_np.T
    
    # 获取维度信息
    batch_size, in_size = a_np.shape
    _, out_size = b_np.shape
    
    # 初始化输出数组
    output_np = np.zeros((batch_size, out_size), dtype=np.float32)
    
    # 调用C实现的矩阵乘法
    c_approx.custom_matmul(a_np, b_np, batch_size, in_size, out_size, output_np)
    
    # 转回TF张量并固定形状
    output = tf.convert_to_tensor(output_np)
    output.set_shape((batch_size, out_size))
    
    # 反向传播：复用TF原生matmul的梯度逻辑
    def grad_fn(dy):
        # 还原原始输入的转置状态，匹配TF原生梯度计算逻辑
        original_a = tf.transpose(a) if transpose_a else a
        original_b = tf.transpose(b) if transpose_b else b
        
        # 计算输入a和b的梯度
        grad_a = tf.matmul(dy, tf.transpose(original_b))
        grad_b = tf.matmul(tf.transpose(original_a), dy)
        
        # 对应前向的转置操作，调整梯度形状
        if transpose_a:
            grad_a = tf.transpose(grad_a)
        if transpose_b:
            grad_b = tf.transpose(grad_b)
        
        # transpose_a/b为静态参数，梯度返回None
        return grad_a, grad_b, None, None
    
    return output, grad_fn

2. 替换Dense层的矩阵乘法 #

自定义Dense层，在call方法中使用上述自定义matmul：

class CustomDense(tf.keras.layers.Dense):
    def call(self, inputs):
        outputs = custom_matmul(inputs, self.kernel)
        if self.use_bias:
            outputs = tf.nn.bias_add(outputs, self.bias)
        outputs = self.activation(outputs)
        return outputs

二、卷积运算（替换tf.nn.conv2d） #

同理，用tf.custom_gradient封装自定义卷积，反向复用TF原生卷积的梯度：

@tf.custom_gradient
def custom_conv2d(inputs, filters, strides, padding):
    # 前向传播：调用C代码
    inputs_np = inputs.numpy()
    filters_np = filters.numpy()
    
    # 获取维度信息（匹配TF默认格式）
    batch, h, w, in_ch = inputs_np.shape
    filter_h, filter_w, _, out_ch = filters_np.shape
    
    # 计算输出维度（严格匹配TF的padding逻辑）
    if padding == "SAME":
        out_h = (h + strides[0] - 1) // strides[0]
        out_w = (w + strides[1] - 1) // strides[1]
    else:
        out_h = (h - filter_h) // strides[0] + 1
        out_w = (w - filter_w) // strides[1] + 1
    
    # 初始化输出数组
    output_np = np.zeros((batch, out_h, out_w, out_ch), dtype=np.float32)
    
    # 调用C实现的卷积运算（参数需匹配你的C代码接口）
    c_approx.custom_conv2d(inputs_np, filters_np,
                          batch, h, w, in_ch,
                          filter_h, filter_w, out_ch,
                          strides[0], strides[1], padding.encode('utf-8'),
                          output_np)
    
    # 转回TF张量并固定形状
    output = tf.convert_to_tensor(output_np)
    output.set_shape((batch, out_h, out_w, out_ch))
    
    # 反向传播：复用TF原生卷积的梯度逻辑
    def grad_fn(dy):
        # 用TF原生API计算输入和卷积核的梯度
        grad_inputs = tf.nn.conv2d_backprop_input(
            input_sizes=tf.shape(inputs),
            filter=filters,
            out_backprop=dy,
            strides=strides,
            padding=padding
        )
        grad_filters = tf.nn.conv2d_backprop_filter(
            input=inputs,
            filter_sizes=tf.shape(filters),
            out_backprop=dy,
            strides=strides,
            padding=padding
        )
        return grad_inputs, grad_filters, None, None
    
    return output, grad_fn

替换Conv2D层 #

自定义Conv2D层，使用上述自定义卷积：

class CustomConv2D(tf.keras.layers.Conv2D):
    def call(self, inputs):
        outputs = custom_conv2d(
            inputs,
            self.kernel,
            strides=self.strides,
            padding=self.padding.upper()
        )
        if self.use_bias:
            outputs = tf.nn.bias_add(outputs, self.bias, data_format=self.data_format)
        outputs = self.activation(outputs)
        return outputs

三、失效原因说明 #

你之前使用tf.py_function的问题在于，它是脱离TensorFlow计算图的Python操作，TF无法追踪其梯度信息，导致反向传播时梯度断裂。而tf.custom_gradient允许手动定义梯度逻辑，这里直接复用TF原生运算的梯度实现，既保证了梯度正确性，又替换了前向计算。

四、训练LeNet-5模型 #

直接替换原模型中的Dense和Conv2D为自定义层即可：

def build_lenet5():
    model = tf.keras.Sequential([
        CustomConv2D(6, kernel_size=(5,5), activation='relu', input_shape=(28,28,1)),
        tf.keras.layers.MaxPooling2D(pool_size=(2,2)),
        CustomConv2D(16, kernel_size=(5,5), activation='relu'),
        tf.keras.layers.MaxPooling2D(pool_size=(2,2)),
        tf.keras.layers.Flatten(),
        CustomDense(120, activation='relu'),
        CustomDense(84, activation='relu'),
        CustomDense(10, activation='softmax')
    ])
    return model

model = build_lenet5()
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# 正常执行训练流程
model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

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