TensorFlow Keras：自定义含多激活函数的Layer报错求助

我尝试构建一个可在keras.models.Sequential中使用的自定义keras.layer.Layer，要在单个Layer中集成多个激活函数。思路是用多个Dense层，每个针对不同数量的单元使用不同激活函数，比如8个单元的层，一半用relu，一半用linear，最后合并两个Dense层的输出。

自定义层代码：

class Multi_Activation_Layer(tf.keras.layers.Layer):

    def __init__(self, units, activations = None, **kwargs):
        self.units = units
        self.activations = activations
        num_fns = len(self.activations)
        remainder = units % num_fns
        adj_units = units - remainder
        units_per_fn = adj_units / num_fns
        self.activations_dict = {activations[-1]: {"Units": units_per_fn + remainder}}

        for fn in self.activations[:-1]:
            self.activations_dict[fn] = {"Units": units_per_fn}

        super(Multi_Activation_Layer, self).__init__(**kwargs)

    def build(self, input_shape):
        
        for activation, dict in self.activations_dict.items():
            dense = tf.keras.layers.Dense(units = dict["Units"], activation = activation)
            dense.build(input_shape)
            self.activations_dict[activation][f"Dense"] = dense

    def call(self, input_data):
        output_data_list = []

        for dict in self.activations_dict.values():
            output = dict["Dense"].call(input_data)
            output_data_list.append(output)

        output_data = tf.concat(output_data_list, axis = 0)
        return output_data
    
    def compute_output_shape(self, input_shape): 
        return (input_shape[0], self.units)

将其加入Sequential模型：

model = tf.keras.models.Sequential([
    Multi_Activation_Layer(8, activations = ["relu", "linear"]),
    tf.keras.layers.Dense(1, activation = "linear")
])

loss_fn = tf.keras.losses.MeanAbsolutePercentageError()

adm_opt = tf.keras.optimizers.Adam(learning_rate = 0.001)

model.compile(
    optimizer = adm_opt,
    loss = loss_fn,
    metrics = [tf.keras.metrics.MeanSquaredError(name = "mse"), tf.keras.metrics.MeanAbsoluteError(name = "mae")]
)

训练时执行model.fit(Xt_train, y_train, epochs = 50)，出现错误：

TypeError                                 Traceback (most recent call last)
c:\Users\XXXXXXX\Documents\XXXXXX\XXXXXXX\XXXXXX\XXXXXXX\XXXXXXX\XXXXXXX Data Preparation.ipynb Cell 225 in ()
----> 1 model.fit(Xt_train, y_train, epochs = 50)

File ~\AppData\Roaming\Python\Python39\site-packages\keras\src\utils\traceback_utils.py:70, in filter_traceback..error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~\AppData\Local\Temp\__autograph_generated_fileheb_wsb5.py:15, in outer_factory..inner_factory..tf__train_function(iterator)
     13 try:
     14     do_return = True
---> 15     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16 except:
     17     do_return = False

File ~\AppData\Local\Temp\__autograph_generated_filedzpps6u7.py:27, in outer_factory..inner_factory..tf__call(self, input_data)
     25 dict = ag__.Undefined('dict')
     26 ag__.for_stmt(ag__.converted_call(ag__.ld(self).activations_dict.values, (), None, fscope), None, loop_body, get_state, set_state, ('dict',), {'iterate_names': 'dict'})
---> 27 output_data = ag__.converted_call(ag__.ld(tf).concat, (ag__.ld(output_data_list),), dict(axis=1), fscope)
     28 try:
     29     do_return = True

TypeError: in user code:

    File "C:\Users\XXXXXX\AppData\Roaming\Python\Python39\site-packages\keras\src\engine\training.py", line 1338, in train_function  *
        return step_function(self, iterator)
    File "C:\Users\XXXXXXX\AppData\Roaming\Python\Python39\site-packages\keras\src\engine\training.py", line 1322, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\XXXXXX\AppData\Roaming\Python\Python39\site-packages\keras\src\engine\training.py", line 1303, in run_step  **
        outputs = model.train_step(data)
    File "C:\Users\XXXXXXX\AppData\Roaming\Python\Python39\site-packages\keras\src\engine\training.py", line 1080, in train_step
        y_pred = self(x, training=True)
    File "C:\Users\XXXXXXXX\AppData\Roaming\Python\Python39\site-packages\keras\src\utils\traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "C:\Users\XXXXXXXX\AppData\Local\Temp\__autograph_generated_filedzpps6u7.py", line 27, in tf__call
        output_data = ag__.converted_call(ag__.ld(tf).concat, (ag__.ld(output_data_list),), dict(axis=1), fscope)

    TypeError: Exception encountered when calling layer 'multi__activation__layer_1' (type Multi_Activation_Layer).
    
    in user code:
    
        File "C:\Users\XXXXXXX\AppData\Local\Temp\ipykernel_19304\3036281980.py", line 42, in call  *
            output_data = tf.concat(output_data_list, axis = 1)
    
        TypeError: '_DictWrapper' object is not callable
    
    
    Call arguments received by layer 'multi__activation__layer_1' (type Multi_Activation_Layer):
      • input_data=tf.Tensor(shape=(None, 37), dtype=float32)

问题根源 #

1. 变量名冲突：循环中使用dict作为变量名，覆盖了Python内置的dict()函数，导致Autograph转换时调用dict(axis=1)出错（此时dict是循环中的字典实例，而非内置函数）。
2. 维度拼接错误：原代码用axis=0拼接，会把样本维度合并（比如batch=32的话，输出会变成64行），正确应该用axis=1拼接特征维度，保证输出形状为(batch_size, units)。
3. 单元数类型错误：units_per_fn = adj_units / num_fns得到浮点数，Dense层的units参数需要整数，应使用整除//。
4. 子层注册问题：将Dense层存在字典中，Keras无法正确跟踪可训练参数，应将子层作为自定义层的属性存储，确保参数被正确注册。

修改后的代码 #

class MultiActivationLayer(tf.keras.layers.Layer):
    def __init__(self, units, activations=None, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.activations = activations or ["relu"]
        num_fns = len(self.activations)
        
        # 计算每个激活函数对应的单元数，确保为整数
        remainder = units % num_fns
        adj_units = units - remainder
        units_per_fn = adj_units // num_fns
        
        # 存储每个激活对应的单元数和Dense层
        self.dense_layers = []
        # 给最后一个激活分配剩余单元
        for i, act in enumerate(self.activations):
            current_units = units_per_fn + (remainder if i == num_fns -1 else 0)
            self.dense_layers.append(tf.keras.layers.Dense(current_units, activation=act))

    def build(self, input_shape):
        # 直接调用子层的build，Keras会自动跟踪参数
        for dense in self.dense_layers:
            dense.build(input_shape)
        super().build(input_shape)

    def call(self, input_data):
        outputs = [dense(input_data) for dense in self.dense_layers]
        # 在特征维度拼接输出
        return tf.concat(outputs, axis=1)
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.units)

使用示例 #

model = tf.keras.models.Sequential([
    MultiActivationLayer(8, activations=["relu", "linear"]),
    tf.keras.layers.Dense(1, activation="linear")
])

loss_fn = tf.keras.losses.MeanAbsolutePercentageError()
adm_opt = tf.keras.optimizers.Adam(learning_rate=0.001)

model.compile(
    optimizer=adm_opt,
    loss=loss_fn,
    metrics=[tf.keras.metrics.MeanSquaredError(name="mse"), tf.keras.metrics.MeanAbsoluteError(name="mae")]
)

# 训练模型
model.fit(Xt_train, y_train, epochs=50)

关键修改说明 #

• 避免使用dict作为变量名，改用列表存储Dense层，彻底解决变量名冲突问题。
• 用整数除法//替换浮点数除法/，确保Dense层的单元数为整数类型。
• 将子Dense层存储为自定义层的列表属性self.dense_layers，让Keras能正确识别并跟踪可训练参数。
• 拼接维度改为axis=1，保证输出形状符合(batch_size, units)的预期，避免后续层维度不匹配。

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