如何在TensorFlow Keras预训练EfficientNetB0中用x*sigmoid(x)替换Swish层？

1. 自定义替代激活层 #

首先实现和Swish原始定义一致的x*sigmoid(x)激活函数，封装为Keras层：

import tensorflow as tf
from tensorflow.keras.layers import Layer

class SwishReplacement(Layer):
    def call(self, x):
        return x * tf.sigmoid(x)

2. 递归遍历模型替换激活 #

EfficientNetB0中的Swish激活可能是单独的Activation层，也可能绑定在卷积、全连接层的activation参数中，因此需要递归遍历所有层完成替换，同时保留原模型权重：

def replace_swish_with_custom(model):
    def _process_layer(layer):
        # 处理嵌套子模型
        if isinstance(layer, tf.keras.Model):
            return tf.keras.Model(inputs=layer.inputs, outputs=_process_layer(layer.output))
        # 替换单独的Swish激活层
        elif isinstance(layer, tf.keras.layers.Activation) and layer.activation.__name__ == 'swish':
            return SwishReplacement(name=f"{layer.name}_replaced")(layer.input)
        # 替换绑定在卷积/全连接层的Swish激活
        elif hasattr(layer, 'activation') and layer.activation is not None and layer.activation.__name__ == 'swish':
            layer.activation = None
            x = layer(layer.input)
            return SwishReplacement(name=f"{layer.name}_swish_replaced")(x)
        # 其他层直接返回原输出
        else:
            return layer(layer.input)
    
    # 构建新模型并复制原权重
    new_model = tf.keras.Model(inputs=model.inputs, outputs=_process_layer(model.output))
    new_model.set_weights(model.get_weights())
    return new_model

3. 实际使用示例 #

# 加载预训练EfficientNetB0
base_model = tf.keras.applications.EfficientNetB0(weights='imagenet', include_top=False)
# 替换Swish激活
new_model = replace_swish_with_custom(base_model)
# 转换为TF-Lite模型
converter = tf.lite.TFLiteConverter.from_keras_model(new_model)
tflite_model = converter.convert()
# 保存模型
with open('efficientnetb0_replaced.tflite', 'wb') as f:
    f.write(tflite_model)

注意事项 #

• 替换后可对比原模型与新模型的输出，确认x*sigmoid(x)和原生Swish的输出一致（二者数学定义完全相同）
• 递归遍历确保处理EfficientNet中所有嵌套的子模块结构

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