如何在Keras中为前向、反向传播配置不同的激活函数

实现逻辑 #

Keras原生支持该操作，核心通过tf.custom_gradient装饰器拆分前向、反向传播逻辑即可：

• 前向传播：调用你自己的不可微分激活函数计算输出
• 反向传播：返回ReLU对应的梯度值回传

完整实现代码 #

首先定义自定义激活函数：

import tensorflow as tf

# 替换为你自己的不可微分激活函数逻辑
def your_non_differentiable_activation(x):
    # 示例为二值化激活，可按需修改为你的逻辑
    return tf.cast(x > 0, dtype=x.dtype)

@tf.custom_gradient
def forward_custom_backward_relu(x):
    # 前向计算逻辑
    forward_output = your_non_differentiable_activation(x)
    # 定义反向梯度计算逻辑
    def grad(dy):
        # ReLU梯度规则：输入大于0时回传上游梯度，否则回传0
        return dy * tf.cast(x > 0, dtype=x.dtype)
    return forward_output, grad

修改你原模型中的激活函数配置即可：

mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

test_images_discrete = test_images

train_images = train_images / 255.0
test_images = test_images / 255.0

# 替换所有层的activation参数为自定义的激活函数
model = tf.keras.Sequential([
  tf.keras.layers.Conv2D(filters=16, kernel_size=5, strides=2, padding="valid",
                         activation=forward_custom_backward_relu, input_shape=(28,28,1), use_bias=True),
  tf.keras.layers.Conv2D(filters=32, kernel_size=3, strides=2, padding="valid",
                         activation=forward_custom_backward_relu, use_bias=True),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(32, activation=forward_custom_backward_relu),
  tf.keras.layers.Dense(10)
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=["accuracy"])
model.fit(train_images,
          train_labels,
          epochs=5,
          validation_data=(test_images, test_labels))

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